Skip to content

Python API

These entries are generated from public Python objects and their docstrings. The hand-written reference explains configuration choices and runtime behavior.

For constructor calls and working extension examples, see the plugin guide.

Package exports

Public package exports for polarism.

BoundaryCondition

Apply boundary actions around each solver step.

before_step_action()

Return the boundary potential term.

after_step_action(psi)

Apply the boundary action to psi.

Config dataclass

Store the full simulation configuration.

AbstractSolver

Bases: ABC

Define the interface for solver backends.

step(potential, pump, reservoir, boundary_condition, state) abstractmethod

Advance the fields from state.t to state.t + dt.

pump(t) returns the pump array at time t. Evaluate it at the integration stage times; the caller updates state.t after the step.

advance(potential, pump, reservoir, boundary_condition, state, start_step, n_steps)

Advance several steps, preserving controller stage times.

RK4CudaSolver

Bases: AbstractSolver

Generic fused RK4 CUDA solver with overridable kernel and launch hooks.

advance(potential, pump, reservoir, boundary_condition, state, start_step, n_steps)

Run a chunk with reusable buffers and, when eligible, a CUDA graph.

step(potential, pump, reservoir, boundary_condition, state)

Advance the solver by one time step.

LaserFactory

Build laser objects from config data.

create_laser(laser_config, X, Y, precision='double') staticmethod

Create one or more lasers from the config.

RealTimeVisualization

Plot fields and scalars during a run.

update(t, fields=None, scalars=None, scalar_groups=None)

Update the plots with new fields and scalars.

ResultNode

Describe one computed result.

compute_field_device(**context)

Compute the field on the active backend without evaluating the scalar reduction.

compute_device(**context)

Compute the field and reduced value on the active backend.

compute_cpu(**context)

Compute the field and reduced value on the CPU.

compute_field_cpu(**context)

Compute the field value on the CPU.

compute_reduced_cpu(**context)

Compute the reduced value on the CPU.

ResultProvider

Bases: ABC

Define the interface for result providers.

make_result_nodes() abstractmethod

Build the result nodes exposed by this object.

SimulationController

Set up and run a full simulation.

run()

Run the simulation loop.

SimulationGrid2D

Bases: Protocol

Protocol for 2D simulation grids.

SimulationState

Store the wavefunction and current time.

At most one of _psi and psi_k is authoritative between step boundaries. Both may coexist after materialization and until the next mutation. A leading batch axis is supported when batch_size is given.

psi property writable

Return the real-space wavefunction.

materialize_psi(xp)

Materialize real-space psi from cached spectral state.

invalidate_psi_k()

Clear cached spectral state after external real-space mutation.

get_simulation_state()

Return simulation state.

Results2D dataclass

Describe one 2D field view.

ResultScalar dataclass

Describe one scalar result.

ResultScalarGroup dataclass

Describe a grouped scalar result.

ResultsManager

Send result nodes to each registered visitor.

Visitors that declare data_location = "device" receive backend-native arrays without a CPU conversion. Each node is computed at most once per :meth:step call regardless of how many visitors request it.

When a node is only requested by device visitors, the scalar reduction (reduce_dim_fn) is skipped so GPU stream synchronisation is avoided. Nodes also needed by CPU visitors always compute both field and reduction.

Visitor exceptions are isolated: a visitor that raises is disabled for the rest of the run so that it cannot abort the simulation loop. Visitors that expose an abort(exc) method have it called before being disabled.

add_visitor(visitor)

Register a result visitor.

step(t, **context)

Compute needed nodes once and dispatch to all visitors.

StorageVisitor

Bases: ResultVisitor

Write result nodes to storage backends.

needs(nodes)

Return names of nodes that are flagged for saving.

visit_all(nodes, t, **context)

Store one result step.

finalize()

Flush all storage backends.

VisualizationVisitor

Bases: ResultVisitor

Send result nodes to the live visualizer.

needs(nodes)

Return names of nodes that are exposed for visualization.

visit_all(nodes, **context)

Send one result step to the visualizer.

ResultVisitor

Bases: ABC

Define the interface for result visitors.

needs(nodes)

Return the set of node names this visitor requires.

Override in subclasses to avoid computing unused nodes.

visit_all(nodes, **context) abstractmethod

Visit all result nodes.

load_plugins(root=None)

Import every module under <root>/polarism_plugins and return their names.

Plugin modules register their components with the usual registry decorators (register_laser, register_potential, ...) on import. root defaults to the current working directory.

Configuration dataclasses

Simulation config dataclasses.

GridParameters dataclass

Store grid size and box-length settings.

PhysicsConstants dataclass

Store the physics constants for the model.

BoundaryConditionParameters dataclass

Store boundary absorption settings.

PotentialParameters dataclass

Store potential settings.

LaserParameters dataclass

Store laser settings.

ReservoirParameters dataclass

Store reservoir settings.

SolverParameters dataclass

Store solver settings.

ResultParameters dataclass

Store result-output settings.

ComputeEngineParameters dataclass

Store backend settings.

Config dataclass

Store the full simulation configuration.

Controller and solver extension

Set up and run a full simulation.

run()

Run the simulation loop.

Bases: ABC

Define the interface for solver backends.

step(potential, pump, reservoir, boundary_condition, state) abstractmethod

Advance the fields from state.t to state.t + dt.

pump(t) returns the pump array at time t. Evaluate it at the integration stage times; the caller updates state.t after the step.

advance(potential, pump, reservoir, boundary_condition, state, start_step, n_steps)

Advance several steps, preserving controller stage times.

Register solver.

Model component extension

The laser, reservoir, solver, grid, and absorption guides explain these bases and registration decorators in context.

Bases: ABC

Define the interface for laser profiles.

power_coefficient(t)

Return the scalar time coefficient of the cached spatial envelope.

get_power(X, Y, t)

Return the full pump field at time t; callers must treat it as read-only.

Register laser.

Bases: ABC

Define the interface for reservoir models.

step(dt, psi_mid, pump) abstractmethod

Advance the reservoir by one time step.

get_reservoir_density() abstractmethod

Return the total reservoir density.

make_result_nodes() abstractmethod

Build the result nodes exposed by this object.

get_state() abstractmethod

Return the current state.

set_state(state) abstractmethod

Set the current state.

get_derivatives(psi, pump, state, rho=None) abstractmethod

Return the reservoir time derivatives.

get_active_density(state)

Return the active reservoir density.

get_inactive_density(state)

Return the inactive reservoir density, or zero if the model has none.

Register reservoir.

Register potential.

Register grid.

Register absorption.

Batched state and reservoir API

SimulationState(grid, eps=1e-3, precision="double", *, mode="legacy_positive_uniform", k_cutoff_um=None, seed=None, batch_size=None) creates identical seeded initial fields along the batch axis. create_reservoir(reservoir_config, physics, grid, precision="double", *, batch_size=None) creates matching zero density fields. Built-in reservoirs support reset(batch_size=None) to clear and resize their density fields. See Batched simulations for a working pattern.