rubix.core package

Contents

rubix.core package#

Submodules#

rubix.core.cosmology module#

rubix.core.cosmology.get_cosmology(config: dict) → BaseCosmology[source]#

Build the requested cosmology object described by config.

Parameters:

config (dict) – Configuration dictionary containing a cosmology entry with name plus optional args for CUSTOM.

Returns:

The selected cosmology implementation.

Return type:

RubixCosmology

Raises:

ValueError – When config["cosmology"]["name"] is not supported.

Example

>>> config = {
...     ...
...     "cosmology":
...         {"name": "PLANCK15"},
...     ...
... }

rubix.core.data module#

class rubix.core.data.Galaxy(redshift: Any | None = None, center: Any | None = None, halfmassrad_stars: Any | None = None)[source]#

Bases: object

Dataclass for storing galaxy metadata.

redshift#

Redshift of the galaxy.

Type:

Optional[Any]

center#

Center coordinates of the galaxy.

Type:

Optional[Any]

halfmassrad_stars#

Half mass radius of the stars.

Type:

Optional[Any]

class rubix.core.data.GasData(coords: Any | None = None, velocity: Any | None = None, mass: Any | None = None, density: Any | None = None, internal_energy: Any | None = None, metallicity: Any | None = None, metals: Any | None = None, sfr: Any | None = None, electron_abundance: Any | None = None, pixel_assignment: Any | None = None, spatial_bin_edges: Any | None = None, mask: Any | None = None, spectra: Any | None = None, datacube: Any | None = None)[source]#

Bases: object

Dataclass for storing the gas component data.

coords#

Coordinates of the gas particles.

Type:

Optional[Any]

velocity#

Velocities of the gas particles.

Type:

Optional[Any]

mass#

Mass of the gas particles.

Type:

Optional[Any]

density#

Density of the gas particles.

Type:

Optional[Any]

internal_energy#

Internal energy values.

Type:

Optional[Any]

metallicity#

Metallicity of the gas particles.

Type:

Optional[Any]

metals#

Metal tracers attached to the particles.

Type:

Optional[Any]

sfr#

Star formation rate values.

Type:

Optional[Any]

electron_abundance#

Electron abundance values.

Type:

Optional[Any]

pixel_assignment#

Pixel assignment in the IFU grid.

Type:

Optional[Any]

spatial_bin_edges#

Spatial bin edges.

Type:

Optional[Any]

mask#

Mask data for the gas.

Type:

Optional[Any]

spectra#

Spectra for each gas particle.

Type:

Optional[Any]

datacube#

IFU datacube of the gas component.

Type:

Optional[Any]

metals: Any | None = None#
class rubix.core.data.RubixData(galaxy: Galaxy | None = None, stars: StarsData | None = None, gas: GasData | None = None)[source]#

Bases: object

Dataclass for storing Rubix data. The RubixData object contains the galaxy, stars, and gas data.

galaxy#

Galaxy metadata.

Type:

Optional[Galaxy]

stars#

Stellar part data.

Type:

Optional[StarsData]

gas#

Gas part data.

Type:

Optional[GasData]

class rubix.core.data.StarsData(coords: Any | None = None, velocity: Any | None = None, mass: Any | None = None, metallicity: Any | None = None, age: Any | None = None, pixel_assignment: Any | None = None, spatial_bin_edges: Any | None = None, mask: Any | None = None, extinction: Any | None = None, spectra: Any | None = None, datacube: Any | None = None)[source]#

Bases: object

Dataclass for storing the stellar component data.

coords#

Coordinates of the stars.

Type:

Optional[Any]

velocity#

Velocities of the stars.

Type:

Optional[Any]

mass#

Mass of the stars.

Type:

Optional[Any]

metallicity#

Metallicity of the stars.

Type:

Optional[Any]

age#

Age of the stars.

Type:

Optional[Any]

pixel_assignment#

Pixel assignment in the IFU grid.

Type:

Optional[Any]

spatial_bin_edges#

Spatial bin edges.

Type:

Optional[Any]

mask#

Mask for the stars.

Type:

Optional[Any]

extinction#

Extinction per particle.

Type:

Optional[Any]

spectra#

Spectra for each stellar particle.

Type:

Optional[Any]

datacube#

IFU datacube of the stellar component.

Type:

Optional[Any]

extinction: Any | None = None#
rubix.core.data.convert_to_rubix(config: dict | str)[source]#

This function converts the data to Rubix format. The data can be loaded from an API or from a file, is then converted to Rubix format and saved to a file (hdf5 format). This ensures that the Rubix pipeline depends not on the simulation data format and basically can hndle any data. If the file already exists, the conversion is skipped.

Parameters:

config (Union[dict, str]) – Configuration dict or path to a YAML file describing the conversion.

Returns:

The output directory where rubix_galaxy.h5 is written.

Return type:

str

Raises:

ValueError – When config['data']['name'] is unsupported.

Example

>>> import os
>>> from rubix.core.data import convert_to_rubix

>>> # Define the configuration (example configuration)
>>> config = {
...    "logger": {
...        "log_level": "DEBUG",
...        "log_file_path": None,
...        "format": "%(asctime)s - %(name)s - %(levelname)s - %(message)s",
...    },
...    "data": {
...        "name": "IllustrisAPI",
...        "args": {
...            "api_key": os.environ.get("ILLUSTRIS_API_KEY"),
...            "particle_type": ["stars","gas"],
...            "simulation": "TNG50-1",
...            "snapshot": 99,
...            "save_data_path": "data",
...        },
...        "load_galaxy_args": {
...            "id": 12,
...            "reuse": True,
...        },
...        "subset": {
...            "use_subset": True,
...            "subset_size": 1000,
...        },
...    },
...    "simulation": {
...        "name": "IllustrisTNG",
...        "args": {
...            "path": "data/galaxy-id-12.hdf5",
...        },
...    },
...    "output_path": "output",
... }

>>> # Convert the data to Rubix format
>>> convert_to_rubix(config)
rubix.core.data.get_reshape_data(config: dict | str) → Callable[source]#

Returns a function to reshape the data

Maps the reshape_array function to the input data dictionary.

Parameters:

config (Union[dict, str]) – Configuration dict or path to the YAML file describing the conversion.

Returns:

Function that reshapes a RubixData instance.

Return type:

Callable[[RubixData], RubixData]

Example

>>> from rubix.core.data import get_reshape_data
>>> reshape_data = get_reshape_data(config)
>>> rubixdata = reshape_data(rubixdata)
rubix.core.data.get_rubix_data(config: dict | str) → RubixData[source]#

Returns the Rubix data. First the function converts the data to Rubix format (convert_to_rubix(config)) and then prepares the input data (prepare_input(config)).

Parameters:

config (Union[dict, str]) – Configuration dict or YAML file path for conversion.

Returns:

RubixData object containing the galaxy, stars, and gas data.

Return type:

RubixData

rubix.core.data.prepare_input(config: dict | str) → RubixData[source]#

Load the converted Rubix dataset into Python objects.

Parameters:

config (Union[dict, str]) – Configuration dict or path describing the conversion.

Returns:

The dataset containing galaxy, stars, and gas objects.

Return type:

RubixData

Raises:

ValueError – When subset mode is enabled but neither stars nor gas coordinates exist.

Example

>>> import os
>>> from rubix.core.data import convert_to_rubix, prepare_input

>>> # Define the configuration (example configuration)
>>> config = {
>>>            ...
>>>           }

>>> # Convert the data to Rubix format
>>> convert_to_rubix(config)

>>> # Prepare the input data
>>> rubixdata = prepare_input(config)
>>> # Access the galaxy data, e.g. the stellar coordintates
>>> rubixdata.stars.coords
rubix.core.data.process_attributes(obj: StarsData | GasData, logger: Logger) → None[source]#

Process the attributes of the given object and reshape them if they are arrays.

rubix.core.data.reshape_array(arr: Array) → Array[source]#

Reshape an array so it can be sharded across devices.

The function reshapes an array of shape (n_particles, n_features) to an array of shape (n_gpus, particles_per_gpu, n_features). Padding with zero is added if necessary to ensure that the number of particles per GPU is the same for all GPUs.

Parameters:

arr (jax.Array) – Array of shape (n_particles, n_features) that should be spread over devices.

Returns:

Array shaped as (n_gpus, particles_per_gpu, ...) suitable for JAX parallelism.

Return type:

jax.Array

rubix.core.ifu module#

rubix.core.ifu.get_calculate_datacube_particlewise(config: dict) → Callable[source]#

Prepare a per-particle datacube builder for the star component.

The returned callable performs an SSP lookup, scales by mass, applies the Doppler shift, resamples onto the telescope wavelength grid, and aggregates the flux into spatial pixels.

First, it looks up the SSP spectrum for each star based on its age and metallicity, scales it by the star’s mass, applies a Doppler shift based on the star’s velocity, resamples the spectrum onto the telescope’s wavelength grid, and finally accumulates the resulting spectra into the appropriate pixels of the datacube.

Parameters:

config (dict) – Configuration dictionary containing telescope and galaxy parameters.

Returns:

Function that computes stars.datacube.

Return type:

Callable[[RubixData], RubixData]

rubix.core.ifu.get_calculate_dusty_datacube_particlewise(config: dict) → Callable[source]#

Prepare a dusty per-particle datacube builder for the star component.

The returned callable is similar to get_calculate_datacube_particlewise() but applies wavelength-dependent extinction using the configured dust model.

First, it looks up the SSP spectrum for each star based on its age and metallicity, scales it by the star’s mass, applies a Doppler shift based on the star’s velocity, resamples the spectrum onto the telescope’s wavelength grid, and finally accumulates the resulting spectra into the appropriate pixels of the datacube.

Parameters:

config (dict) – Configuration dictionary containing telescope and galaxy parameters as well as ssp.dust settings.

Returns:

Function that computes stars.datacube with extinction.

Return type:

Callable[[RubixData], RubixData]

rubix.core.lsf module#

rubix.core.lsf.get_convolve_lsf(config: dict) → Callable[[RubixData], RubixData][source]#

Create the LSF convolution function described by config.

Parameters:

config (dict) – Configuration dict that must include telescope.lsf.

Returns:

Function that convolves Rubix data.

Return type:

Callable[[RubixData], RubixData]

Raises:

ValueError – When the telescope LSF configuration or sigma is missing.

Example

>>> config = {
...     ...
...     "telescope": {
...         "name": "MUSE",
...         "psf": {"name": "gaussian", "size": 5, "sigma": 0.6},
...         "lsf": {"sigma": 0.5},
...         "noise": {"signal_to_noise": 1,"noise_distribution": "normal"},
...    },
...     ...
... }

>>> from rubix.core.lsf import get_convolve_lsf
>>> convolve_lsf = get_convolve_lsf(config)
>>> rubixdata = convolve_lsf(rubixdata)

rubix.core.noise module#

rubix.core.noise.get_apply_noise(config: dict) → Callable[[RubixData], RubixData][source]#

Build the noise application function described by config.

Parameters:

config (dict) – Configuration dict that includes telescope.noise.

Returns:

Function that adds noise to data.

Return type:

Callable[[RubixData], RubixData]

Raises:

ValueError – When required noise configuration keys are missing.

Example

>>> config = {
...     ...
...     "telescope": {
...         "name": "MUSE",
...         "psf": {"name": "gaussian", "size": 5, "sigma": 0.6},
...         "lsf": {"sigma": 0.5},
...         "noise": {"signal_to_noise": 1,"noise_distribution": "normal"},
...    },
...     ...
... }

>>> from rubix.core.noise import get_apply_noise
>>> apply_noise = get_apply_noise(config)
>>> rubixdata = apply_noise(rubixdata)

rubix.core.pipeline module#

class rubix.core.pipeline.RubixPipeline(user_config: dict | str)[source]#

Bases: object

Builds and executes the Rubix data processing pipeline with the provided configuration.

Parameters:

user_config (Union[dict, str]) – Parsed configuration dictionary or path to a configuration file.

Example

>>> from rubix.core.pipeline import RubixPipeline
>>> config = "path/to/config.yml"
>>> target_datacube = ...  # Load or define your target datacube here
>>> pipe = RubixPipeline(config)
>>> inputdata = pipe.prepare_data()
>>> final_datacube = pipe.run_sharded(inputdata)
>>> ssp_model = pipe.ssp
>>> telescope = pipe.telescope
>>> loss_value = pipe.loss(inputdata, target_datacube)
>>> gradient_data = pipe.gradient(inputdata, target_datacube)
gradient(rubixdata: RubixData, targetdata: Array) → RubixData[source]#

Compute the gradient of the loss with respect to rubixdata.

Parameters:

  • rubixdata (RubixData) – Pytree describing the current pipeline input data.

  • targetdata (jnp.ndarray) – Target datacube used to compute the loss.

Returns:

Gradient pytree that matches the structure of rubixdata.

Return type:

RubixData

loss(rubixdata: RubixData, targetdata: Array) → Array[source]#

Compute the mean squared error between pipeline output and target.

Parameters:

  • rubixdata (RubixData) – Input data passed to run().

  • targetdata (jnp.ndarray) – Target datacube used for comparison.

Returns:

Scalar mean squared error value.

Return type:

jnp.ndarray

prepare_data()[source]#

Prepares and loads the data for the pipeline.

Returns:

‘coords’, ‘velocities’, ‘mass’, ‘age’, and ‘metallicity’ under stars and gas.

Return type:

Object containing particle data with attributes such as

run_sharded(inputdata: RubixData, devices: Sequence[Any] | None = None) → Array[source]#

Run the compiled pipeline across devices by sharding the particle data.

It splits the particle arrays under stars and gas into shards,

runs the compiled pipeline on each shard, and then combines the resulting datacubes.

Note:

This is the recommended method to run the pipeline in parallel at the moment.

Args:
inputdata (RubixData):

Output of prepare_data(). Contains star and gas particles.

devices (Optional[Sequence[Any]], optional):

Devices to use for shard_map(). These should be jax.Device instances. Defaults to jax.devices().

Returns:

jnp.ndarray: Sharded pipeline output aggregated across devices.

rubix.core.psf module#

rubix.core.psf.get_convolve_psf(config: dict) → Callable[source]#

Return a callable that applies the configured PSF kernel.

Parameters:

config (dict) – Pipeline configuration that must include telescope settings. The telescope.psf block requires name and, when using the Gaussian kernel, size and sigma fields to define the kernel dimensions and width.

Returns:

Callable that convolves the stars

datacube with the generated PSF kernel.

Return type:

Callable[[RubixData], RubixData]

Raises:

ValueError – When the PSF settings are missing or reference an unknown kernel type.

Example

>>> config = {
...     ...
...     "telescope": {
...         "name": "MUSE",
...         "psf": {"name": "gaussian", "size": 5, "sigma": 0.6},
...         "lsf": {"sigma": 0.5},
...         "noise": {
...             "signal_to_noise": 1,
...             "noise_distribution": "normal",
...         },
...    },
...     ...
... }

>>> from rubix.core.psf import get_convolve_psf
>>> convolve_psf = get_convolve_psf(config)
>>> rubixdata = convolve_psf(rubixdata)

rubix.core.rotation module#

rubix.core.rotation.get_galaxy_rotation(config: dict)[source]#

Return a rotation function configured for the provided galaxy settings.

Parameters:

config (dict) – Pipeline configuration containing galaxy.rotation with either type (face-on, edge-on, matrix) or explicit alpha, beta, gamma angles.

Returns:

Function that applies the requested rotation.

Return type:

Callable[[RubixData], RubixData]

Raises:

ValueError – When the rotation configuration or required fields are invalid or missing.

Example

>>> config = {
...     ...
...     "galaxy": {
...         "dist_z": 0.1,
...         "rotation": {"type": "edge-on"},
...     },
...     ...
... }

>>> from rubix.core.rotation import get_galaxy_rotation
>>> rotate_galaxy = get_galaxy_rotation(config)
>>> rubixdata = rotate_galaxy(rubixdata)

rubix.core.ssp module#

rubix.core.ssp.get_lookup_interpolation(config: dict) → Callable[source]#

Loads the SSP template defined in the configuration and returns the lookup function for the template.

The lookup function is a function that takes in the metallicity and age of a star and returns the spectrum of the star. This is later used to vmap over the stars metallicities and ages, and pmap over multiple GPUs.

Parameters:

config (dict) – Configuration dictionary.

Returns:

Lookup function for the SSP template.

rubix.core.ssp.get_lookup_interpolation_pmap(config: dict) → Callable[source]#

Get the pmap version of the lookup function for the SSP template defined in the configuration.

Parameters:

config (dict) – Configuration dictionary.

Returns:

pmapped lookup function for the SSP template.

rubix.core.ssp.get_lookup_interpolation_vmap(config: dict) → Callable[source]#

This function loads the SSP template defined in the configuration and returns the lookup function for the template, vmapped over the stars metallicities and ages.

Parameters:

config (dict) – Configuration dictionary.

Returns:

vmapped lookup function for the SSP template.

rubix.core.ssp.get_ssp(config: dict) → object[source]#

This function loads the simple stellar population (SSP) template defined in the configuration.

Parameters:

config (dict) – Configuration dictionary.

Returns:

SSP template

Raises:

ValueError – If required keys are missing from the configuration.

rubix.core.telescope module#

rubix.core.telescope.get_filter_particles(config: dict) → Callable[source]#

Return a callable that masks particles outside the aperture.

Parameters:

config (dict) – Configuration dictionary.

Returns:

Function that filters particles.

Return type:

Callable[[RubixData], RubixData]

Example

>>> from rubix.core.telescope import get_filter_particles
>>> filter_particles = get_filter_particles(config)

>>> rubixdata = filter_particles(rubixdata)
rubix.core.telescope.get_spatial_bin_edges(config: dict) → Float[Array, 'n_bins'][source]#

Compute the spatial bin edges that map particles into spaxels.

Parameters:

config (dict) – Configuration dictionary containing telescope and galaxy data.

Returns:

Array of spatial bin edges.

Return type:

Float[Array, n_bins]

rubix.core.telescope.get_spaxel_assignment(config: dict) → Callable[source]#

Return a particles-to-spaxels assignment function.

Parameters:

config (dict) – Configuration dictionary.

Returns:

Function that assigns particles to spaxels.

Return type:

Callable[[RubixData], RubixData]

Raises:

ValueError – If the telescope pixel type is unsupported.

Example

::
>>> from rubix.core.telescope import get_spaxel_assignment
>>> bin_particles = get_spaxel_assignment(config)

>>> rubixdata = bin_particles(rubixdata)

>>> print(rubixdata.stars.pixel_assignment)
>>> print(rubixdata.stars.spatial_bin_edges)
rubix.core.telescope.get_telescope(config: str | dict) → BaseTelescope[source]#

Return the configured telescope instance.

Parameters:

config (Union[str, dict]) – Configuration dictionary or path that includes telescope.name.

Returns:

Concrete telescope object assembled by the factory.

Return type:

BaseTelescope

Raises:

TypeError – If the factory does not return a BaseTelescope.

Example

::
>>> from rubix.core.telescope import get_telescope
>>> config = {
...     "telescope": {"name": "MUSE"},
... }
>>> telescope = get_telescope(config)
>>> print(telescope)

Module contents#

Contents