"The test of a first-rate intelligence is the ability to hold two opposed ideas in the mind at the same time, and still retain the ability to function." ~ F. Scott Fitzgerald

Polarism

Polarism is a Python package for 2-D driven-dissipative Gross-Pitaevskii simulations of exciton-polariton systems. It provides configurable pumps, reservoirs, solvers, and result output for CPU and optional CUDA runs.

_{A ring-shaped pump (the example plugin below) traps a rotating condensate with quantized vortices. Left: density |ψ|², right: phase. 256² grid; 180 frames spanning 300–1200 ps, rendered by examples/readme_media.py.}

Highlights

• Finite-difference RK4, split-step FFT, ETD-RK2, and interaction-picture solver families.
• NumPy execution and optional CuPy GPU execution, including CUDA graph chunks for eligible CUDA solver runs.
• Single, double, and quadratic-double reservoirs; continuous and pulsed pumps.
• HDF5, NPY, and JSON output written on a background thread, so saving frames barely slows GPU runs; ps.load_results reads runs back lazily.
• Checkpoint and restart, including a clean stop on SLURM time-limit signals.
• Parameter sweeps with batched device execution, saved member results, and resume.
• Live visualization and online GPU-rendered animation.
• Registries for custom solvers, model components, output formats, and observables, documented in the plugin guide; a tyro-generated polarism CLI.
• polarism-bench for measuring per-step and chunked GPU execution on your hardware.

Installation

pip install polarism

For CUDA runs, install the extra matching your CUDA environment:

The installed CuPy build must be compatible with the NVIDIA driver and CUDA toolkit. See the installation guide and GPU guide.

For a development checkout:

pip install -e ".[dev]"

Quickstart

This example matches examples/basic_run.py:

import polarism as ps
from polarism.compute_engine import compute_engine

cfg = ps.Config()

cfg.grid.nx = 128
cfg.grid.ny = 128
cfg.grid.lx = 100.0
cfg.grid.ly = 100.0
cfg.grid.grid_type = "periodic"

cfg.reservoir.reservoir_type = "single"
cfg.potential.potential_type = "zero"

cfg.laser.mode = "single"
cfg.laser.laser_type = "continuous-gaussian"
cfg.laser.P0 = 0.02
cfg.laser.sigma_space = 12.0

cfg.solver.method = "rk4-fdm"
cfg.solver.dt = 1e-3
cfg.solver.total_time = 2.0

cfg.result.save_results = True
cfg.result.save_hdf5 = True
cfg.result.output_directory = "quickstart_results"

cfg.compute_engine.use_gpu = False
compute_engine.configure(cfg.compute_engine)

ps.SimulationController(cfg).run()

Open a saved run with ps.load_results(run_dir). Its fields["density"], scalars["norm"], and static["pump"] use stable ASCII names; x and y contain the grid coordinates. HDF5 datasets are read lazily, so fields["density"][::10] loads every tenth frame only. Use the returned object as a context manager for HDF5 output.

You can save this configuration to YAML with cfg.to_yaml("run.yaml"), load it with ps.Config.from_yaml("run.yaml"), or run it from the CLI:

polarism --config run.yaml --grid.nx 256

CLI flags override values in the file. See the configuration guide.

For long cluster jobs, cfg.checkpoint.every_steps writes restartable checkpoints and cfg.checkpoint.resume_from continues from one; see Checkpoint and restart.

For GPU execution, set cfg.compute_engine.use_gpu = True, choose cfg.solver.method = "rk4-cuda", "split-step-fft-cuda", or "ifrk-fft-cuda", and consider cfg.solver.precision = "single" on consumer GPUs. FFT solvers require a periodic grid. Inspect CLI options with polarism --help.

To try Polarism on a GPU without installing CUDA, open the Colab notebook.

Adding your own pump

Pumps, reservoirs, potentials, solvers, grids, and absorption strategies are selected by name from registries. Subclass the base class, register it, and select the name in the config. This ring-shaped pump traps vortices inside the ring:

import polarism as ps
from polarism.laser import register_laser
from polarism.laser.abstract_laser import AbstractLaser


@register_laser("ring")
class RingLaser(AbstractLaser):
    """Continuous Gaussian ring of radius `ring_radius` and width `sigma_space`."""

    def __init__(self, laser_config, X, Y, precision="double"):
        super().__init__(laser_config, X, Y, precision)
        self.radius = laser_config.extra.get("ring_radius", 20.0)
        self.width = laser_config.sigma_space
        self._finalize_spatial_envelope(X, Y)

    def _amplitude(self, t):
        return self.P0

    def _P_space(self, X, Y):
        r = self.xp.sqrt((X - self.x0) ** 2 + (Y - self.y0) ** 2)
        return self.xp.exp(-0.5 * ((r - self.radius) / self.width) ** 2)

    def _P_time(self, t):
        return 1.0


cfg = ps.Config()
cfg.laser.laser_type = "ring"
cfg.laser.extra = {"ring_radius": 20.0}
cfg.laser.sigma_space = 4.0

The class works on both NumPy and CuPy through self.xp. Put it in polarism_plugins/ to have it discovered automatically, including by the polarism CLI. The plugin guide covers the contract for each component type.

Parameter sweeps

Sweep related configs together and read any saved member by index:

import polarism as ps

cfg = ps.Config()
cfg.result.save_results = True
cfg.result.save_hdf5 = True
with ps.Sweep(cfg, {"laser.P0": [0.02, 0.03], "physics.g_C": [0.001, 0.002]}).run() as results:
    density = results.member(0).fields["density"][:]

See the parameter sweeps guide for batching, output, and restart.

Choosing a solver

Check timestep and grid convergence for your model. See the solver reference for compatibility details.

Documentation

• Getting started
• User guide
• Running on GPUs
• Plugin guide
• Reference and Python API
• Physics background

For pulsed pumps, power_definition: pulse_energy makes power or P the integrated dose of one pulse over space and time. The legacy peak_amplitude mode remains available for reproducing older runs. See Potentials and Pumps.

Contributing

See the contributing guide. Run the default regression suite with .venv/bin/pytest -q from a development checkout.

Citation

Cite Polarism with its Zenodo DOI, 10.5281/zenodo.23531116, which always resolves to the latest release; each release also has its own version DOI on Zenodo. The same metadata is in CITATION.cff and GitHub's “Cite this repository” menu.

License

Polarism is licensed under the MIT License.

Acknowledgements

This project is developed with the Exciton-Polariton research group at the Faculty of Physics, University of Warsaw, with support from NCN OPUS grant work and ICM computing infrastructure.