Usage

Installation

To use synthpy, first install it using pip:

(.venv) $ pip install synthPy

Regular usage

To setup the jax environment, you can use the simulator.config.jax_init() function:

simulator.config.jax_init(force_device=None, core_limit=None, extra_info=False, disable_python_multithreading=True, enable_x64=False, debugging=False, jax_updated=True)
Initialises jax and imports

  • sets enviroment variables and runtime configuration

  • sets jax backend (where to run computations) based on autodetection or forced configuration with :param force_device:

Parameters:

  • force_device (str (allowed: "CPU", "GPU", "TPU") or None) – Manually select compute backend.

  • core_limit (int or None) – Limit the number of threads that jax can parallelise across.

  • extra_info (bool (default = False)) – Prints out jax environment info if enabled.

  • disable_python_multithreading (bool (default = True)) – Disables python multithreading to prevent conflict with jax parallelisation in some instances.

  • enable_x64 (bool (default = False)) – Enable 64-bit values in jax (double precision floating point arithmetic).

  • debugging (bool (default = False)) – Enables debug flags, increases runtime.

  • jax_updated (bool (default = True)) – Some flags aren’t available in older versions of jax, set to False if running into issues setting them - should deprecate soon.

Raises:

AssertionError – If jax has already been imported.

Returns:

No return, sets up configuration variables and import without returning values.

Return type:

None

This will setup jax and in the default configuration for this program, it will also configure file paths for imports. There are no required parameters - but there are many options.

After this, there are 3 main parts to running a simulation:

  • Beam initialisation

  • Domain generation

  • Propagation

There is then post-processing to get sensible information and plots out of your raw data, but this will be covered separately.

Beam initialisation

Beam initialisation is essential down to instantiating an object from the simulator.beam.Beam class:

Domain generation

Domain generation is again essential down to instantiating an object from the simulator.domain.ScalarDomain class:

This class inherits from equinox.Module in order to work properly with jax, hence attributes cannot be added, removed or had their shape changed after creation. To work around this you either have to re-create your instance of the simulator.domain.ScalarDomain class or update your current version using equinox.tree_at

Note

We recommend leaving this operation up to us, but if necessary, just re-create the object rather than messing with tree edits - it’s likely safer for jax. (If you want to know why this is an issue then look here at jax tracers: https://docs.jax.dev/en/latest/key-concepts.html#tracing)

We are considering switching this domain creation to a functional approach rather than class based to avoid this very issue and make its creation more versatile, so this may change in the future.

Propagation

For example:

>>> import synthPy
>>> synthPy.simulator.config.jax_init()

There is also a generic default run jupyter notebook in synthPy/examples/notebooks/SynthRayTracer.ipynb

Batching

Initial jax setup is the same as beforere-creating