rubix.core package

rubix.core package#

Submodules#

rubix.core.cosmology module#

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

Get the cosmology from the configuration

Parameters:

config – Configuration dictionary

Returns:

RubixCosmology

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 the galaxy data

Parameters:

  • redshift – Redshift of the galaxy

  • center – Center coordinates of the galaxy

  • halfmassrad_stars – Half mass radius of the stars in the galaxy

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 Gas data

Parameters:

  • coords – Coordinates of the gas particles

  • velocity – Velocities of the gas particles

  • mass – Mass of the gas particles

  • density – Density of the gas particles

  • internal_energy – Internal energy of the gas particles

  • metallicity – Metallicity of the gas particles

  • sfr – Star formation rate of the gas particles

  • electron_abundance – Electron abundance of the gas particles

  • pixel_assignment – Pixel assignment of the gas particles in the IFU grid

  • spatial_bin_edges – Spatial bin edges of the IFU grid

  • mask – Mask for the gas particles

  • spectra – Spectra for each gas particle

  • datacube – IFU datacube for the gas component

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.

Parameters:

  • galaxy – Galaxy object containing the galaxy data

  • stars – StarsData object containing the stars data

  • gas – GasData object containing the gas data

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 stars data

Parameters:

  • coords – Coordinates of the stars

  • velocity – Velocities of the stars

  • mass – Mass of the stars

  • metallicity – Metallicity of the stars

  • age – Age of the stars

  • pixel_assignment – Pixel assignment of the stars in the IFU grid

  • spatial_bin_edges – Spatial bin edges of the IFU grid

  • mask – Mask for the stars

  • spectra – Spectra for each stellar particle

  • datacube – IFU datacube for the stellar component

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 (dict or str) – The configuration for the conversion. This can be a dictionary or a path to a YAML file containing the configuration.

Returns:

The configuration used for the conversion. This can be used to pass the output path to the next step in the pipeline.

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 (dict or str) – The configuration for the conversion. This can be a dictionary or a path to a YAML file containing the configuration.

Returns:

A function to reshape the data.

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 (dict or str) – The configuration for the conversion. This can be a dictionary or a path to a YAML file containing the configuration.

Returns:

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

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

This function prepares the input data for the pipeline. It loads the data from the file and converts it to Rubix format.

Parameters:

config (dict or str) – The configuration for the conversion. This can be a dictionary or a path to a YAML file containing the configuration.

Returns:

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

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]#

Reshapes an array to be compatible with JAX parallelization

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 (jnp.ndarray) – The array to reshape

Returns:

The reshaped array as jnp.ndarray

rubix.core.ifu module#

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

Create a function that calculates the datacube for the stars component of a RubixData object on a per-particle basis. 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:

A function that takes a RubixData object and returns it with

the datacube calculated and added to the stars component.

Return type:

Callable

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

Create a function that calculates the datacube for the stars component of a RubixData object on a per-particle basis. 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:

A function that takes a RubixData object and returns it with

the datacube calculated and added to the stars component.

Return type:

Callable

rubix.core.lsf module#

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

Get the function to convolve with the Line Spread Function (LSF) based on the configuration.

Parameters:

config (dict) – Configuration dictionary.

Returns:

The function to convolve with the LSF.

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[source]#

Get the function to apply noise to the datacube based on the configuration.

Parameters:

config (dict) – Configuration dictionary.

Returns:

The function to apply noise to the datacube.

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

RubixPipeline is responsible for setting up and running the data processing pipeline.

Parameters:

  • user_config (dict or str) – Parsed user configuration for the pipeline.

  • pipeline_config (dict) – Configuration for the pipeline.

  • logger (Logger) – Logger instance for logging messages.

  • ssp (object) – Stellar population synthesis model.

  • telescope (object) – Telescope configuration.

  • data (dict) – Dictionary containing particle data.

  • func (callable) – Compiled pipeline function to process data.

Example

>>> from rubix.core.pipeline import RubixPipeline
>>> config = "path/to/config.yml"
>>> pipe = RubixPipeline(config)
>>> inputdata = pipe.prepare_data()
>>> output = pipe.run(inputdata)
>>> # To run with sharding using jax.shard_map:
>>> final_datacube = pipe.run_sharded(inputdata, shard_size=100000)
>>> ssp_model = pipeline.ssp
>>> telescope = pipeline.telescope
gradient(rubixdata, targetdata)[source]#

This function will calculate the gradient of the pipeline.

loss(rubixdata, targetdata)[source]#

Calculate the mean squared error loss.

Parameters:

  • data (array-like) – The predicted data.

  • target (array-like) – The target data.

Returns:

The mean squared error loss.

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, devices=None)[source]#

Runs the pipeline on sharded input data in parallel using jax.shard_map. It splits the particle arrays (e.g. under stars and gas) into shards, runs the compiled pipeline on each shard, and then combines the resulting datacubes.

This is the recomended method to run the pipeline in parallel at the moment!!!

Parameters:

  • inputdata (object) – Data prepared from the prepare_data method.

  • shard_size (int) – Number of particles per shard.

Returns:

The final datacube combined from all shards.

Return type:

jax.numpy.ndarray

rubix.core.psf module#

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

Get the point spread function (PSF) kernel based on the configuration.

Parameters:

config (dict) – Configuration dictionary.

Returns:

The function to convolve the datacube with the PSF kernel.

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]#

Get the function to rotate the galaxy based on the configuration.

Parameters:

config (dict) – Configuration dictionary.

Returns:

The function to rotate the galaxy.

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

rubix.core.telescope module#

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

Get the function to filter particles outside the aperture.

Parameters:

config (dict) – Configuration dictionary.

Returns:

The filter particles function

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]#

Get the spatial bin edges based on the configuration.

Parameters:

config (dict) – Configuration dictionary.

Returns:

The spatial bin edges.

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

Get the spaxel assignment function based on the configuration.

Parameters:

config (dict) – Configuration dictionary.

Returns:

The spaxel assignment function.

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]#

Get the telescope object based on the configuration.

Parameters:

config (dict) – Configuration dictionary.

Returns:

The telescope object.

Example

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

Module contents#

Contents