borgjaxpm API

class borgjaxpm.BorgJaxPM(sim_box: BoxModel3, a_init: float, a_final: float = 1.0, forcesampling: int = 1, integrator: str = 'bullfrog', num_steps: int = 50, supersampling: int = 1, mult_fix=1, fd: bool = None)

Bases: BaseForwardModel_v3

adjointModel_v3(self: aquila_borg._borg.forward.ForwardModel, input: LibLSS::GeneralIO::details::InputAdjoint) → None

Execute the model in the adjoint gradient direction using the v3 API.

The adjoint gradient direction is also called back-propagation in neural network language.

See getResultAdjointGradient_v3() for a complete example.

Parameters:

input (GInputAdjoint) – input field on mesh

forwardModel_v3(self: aquila_borg._borg.forward.ForwardModel, input: LibLSS::GeneralIO::details::Input) → None

Execute the model in the forward direction using the v3 API

See getResultForward_v3() for a complete example.

Parameters:

input (GInput) – input field on mesh

getInputDescription(self: aquila_borg._borg.forward.ForwardModel) → LibLSS::DataRepresentation::Descriptor

Query the format of the input of the model.

Returns:

a ModelIO descriptor that describes the input of the model

Return type:

RepresentationDescriptor

getOutputDescription(self: aquila_borg._borg.forward.ForwardModel) → LibLSS::DataRepresentation::Descriptor

Query the format of the output of the model.

Returns:

a ModelIO descriptor that describes the output of the model

Return type:

RepresentationDescriptor

getOutputDescription_AG() → RepresentationDescriptor

getParticlePositions() → _Buffer | _SupportsArray[dtype[Any]] | _NestedSequence[_SupportsArray[dtype[Any]]] | complex | bytes | str | _NestedSequence[complex | bytes | str]

getParticleVelocities() → _Buffer | _SupportsArray[dtype[Any]] | _NestedSequence[_SupportsArray[dtype[Any]]] | complex | bytes | str | _NestedSequence[complex | bytes | str]

getResultAdjointGradient_v3(self: aquila_borg._borg.forward.ForwardModel, output: LibLSS::GeneralIO::details::OutputAdjoint) → LibLSS::GeneralIO::details::OutputAdjoint

Retrieve the result of the adjoint gradient evaluation using the v3 API.

Be careful that the provided object is emptied of its content when executing this function. This is done to ensure the consistency of the algorithm (which can request the local representation of the field to be transformed). A valid output is returned which must be used to access the mesh, if the provided object was a temporary.

Parameters:

output (GOutputAdjoint) – an output mesh holder in which to put the result.

Returns:

Returned new holder of the mesh.

Return type:

GOutputAdjoint

Example:

Here is an example on how to use the interface.

Nside = 2
# Setup the default box model
mybox = borg.forward.BoxModel()
# We generated some random data
some_input_data = np.random.randn(*mybox.N)
in_data_rep = borg.modelio.newInputModelIO(mybox, some_input_data)
m.adjointModel_v3(borg.modelio.GInputAdjoint(in_data))

# Setup a tiled array for the output. It is a 1d array here
out_array_rep = borg.modelio.newTiledArray([12 * Nside**2])
# Here no copy of the array exist in python, and no local tiling has been defined
own_array_rep = m.getResultAdjointGradient_v3(borg.modelio.GOutputAdjoint(out_array_rep))
# out_array_rep is an empty representation now. own_array_rep does hold the data
out_data = own_array_rep.disown()
# out_data is now a TiledArray. The representation is empty again and may not be used.

getResultForward_v3(self: aquila_borg._borg.forward.ForwardModel, output: LibLSS::GeneralIO::details::Output) → LibLSS::GeneralIO::details::Output

Retrieve the result of the forward model evaluation using the v3 API.

Be careful that the provided object is emptied of its content when executing this function. This is done to ensure the consistency of the algorithm (which can request the local representation of the field to be transformed). A valid output is returned which must be used to access the mesh, if the provided object was a temporary.

Parameters:

output (GOutput) – an output mesh holder in which to put the result.

Returns:

Returned new holder of the mesh.

Return type:

GOutput

Example:

Here is an example on how to use the interface.

Nside = 2
# Setup the default box model
mybox = borg.forward.BoxModel()
# We generated some random data
some_input_data = np.random.randn(*mybox.N)
in_data_rep = borg.modelio.newInputModelIO(mybox, some_input_data)
m.forwardModel_v3(borg.modelio.GInput(in_data))

# Setup a tiled array for the output. It is a 1d array here
out_array_rep = borg.modelio.newTiledArray([12 * Nside**2])
# Here no copy of the array exist in python, and no local tiling has been defined
own_array_rep = m.getResultForward_v3(borg.modelio.GOutput(out_array_rep))
# out_array_rep is an empty representation now. own_array_rep does hold the data
out_data = own_array_rep.disown()
# out_data is now a TiledArray. The representation is empty again and may not be used.

setAdjointRequired(self: aquila_borg._borg.forward.ForwardModel, arg0: bool) → None

Indicate whether the caller require the adjoint gradient to be computed later.

Parameters:

flag (bool) – if True, tape is recorded to compute adjoint

borgjaxpm.BorgJaxPM_v2

alias of _wrapper_model

borgjaxpm.funcs.prepare_bullfrog_simulation(omega_c, omega_b, num_steps, a_start, a_final, supersampling=1, mult_fix=1, fd=False)

borgjaxpm.funcs.prepare_cola_simulation(omega_c, omega_b, num_steps, a_start, a_final, supersampling=1, lpt_order=2, log_step=True)

Pre-compute COLA coefficients and return a ready-to-call simulation function.

The returned callable has the same signature as the one returned by prepare_leap_frog_simulation, so both can be used interchangeably.

Parameters:

• omega_c (float) – Cold dark matter density parameter Ω_c h².

• omega_b (float) – Baryon density parameter Ω_b h².

• num_steps (int) – Number of PM time steps. COLA typically needs ~10–30× fewer steps than a standard PM simulation for the same large-scale accuracy.

• a_start (float) – Initial scale factor (should match the epoch of ic_field).

• a_final (float) – Final scale factor.

• supersampling (int, optional) – Supersampling factor applied to ic_field before placing particles. Default 1 (no supersampling).

• lpt_order (int, optional) – Order of the LPT reference trajectory (1 or 2). Second order (default) gives better accuracy for the same number of steps.

• log_step (bool, optional) – If True (default) steps are uniform in ln a; otherwise uniform in a.

Returns:

run_fn – run_fn(ic_field, mesh_shape, mesh_shape_force, **kwargs) -> positions

Return type:

callable

Examples

>>> run = prepare_cola_simulation(0.25, 0.05, num_steps=10,
...                               a_start=0.1, a_final=1.0)
>>> positions = run(ic_field, mesh_shape=(128, 128, 128),
...                 mesh_shape_force=(128, 128, 128))

borgjaxpm.funcs.prepare_leap_frog_simulation(omega_c, omega_b, num_steps, a_start, a_final, supersampling=1)

borgjaxpm.funcs.run_simulation_classic(ic_field, omega_c, omega_b, want_vel, mesh_shape, mesh_shape_force, snapshots, a_start, sharding, supersampling)

ic_field: initial condition mesh grid at epoch ‘a_start’, under the forme of the density contrast omega_c: cold dark matter density omega_b: baryon density mesh_shape: shape of the position mesh in lagrangian coordinates mesh_shape_force: shape of the force mesh snaphots: requested snapshots a_start: epoch to start the PM simulation from supersampling: supersampling factor for the initial condition field