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: Array | None = None, center: Array | None = None, halfmassrad_stars: Array | 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: Array | None = None, velocity: Array | None = None, mass: Array | None = None, density: Array | None = None, internal_energy: Array | None = None, metallicity: Array | None = None, sfr: Array | None = None, electron_abundance: Array | None = None, pixel_assignment: Array | None = None, spatial_bin_edges: Array | None = None, mask: Array | None = None, spectra: Array | None = None, datacube: Array | 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

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: Array | None = None, velocity: Array | None = None, mass: Array | None = None, metallicity: Array | None = None, age: Array | None = None, pixel_assignment: Array | None = None, spatial_bin_edges: Array | None = None, mask: Array | None = None, spectra: Array | None = None, datacube: Array | 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

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) → object[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) → object[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.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(config: dict) → Callable[source]#

The function returns the function that calculates the datacube of the stars.

Parameters:

config (dict) – The configuration dictionary

Returns:

The function that calculates the datacube of the stars.

Example

>>> from rubix.core.ifu import get_calculate_datacube
>>> calculate_datacube = get_calculate_datacube(config)

>>> rubixdata = calculate_datacube(rubixdata)
>>> # Access the datacube of the stars
>>> rubixdata.stars.datacube
rubix.core.ifu.get_calculate_spectra(config: dict) → Callable[source]#

The function gets the lookup function that performs the lookup to the SSP model, and parallelizes the funciton across all GPUs.

Parameters:

config (dict) – The configuration dictionary

Returns:

The function that calculates the spectra of the stars.

Example

>>> config = {
...     "ssp": {
...         "template": {
...             "name": "BruzualCharlot2003"
...             },
...          },
...     }

>>> from rubix.core.ifu import get_calculate_spectra
>>> calcultae_spectra = get_calculate_spectra(config)

>>> rubixdata = calcultae_spectra(rubixdata)
>>> # Access the spectra of the stars
>>> rubixdata.stars.spectra
rubix.core.ifu.get_doppler_shift_and_resampling(config: dict) → Callable[source]#

The function doppler shifts the wavelength based on the velocity of the stars and resamples the spectra to the telescope wavelength grid.

Parameters:

config (dict) – The configuration dictionary

Returns:

The function that doppler shifts the wavelength based on the velocity of the stars and resamples the spectra to the telescope wavelength grid.

Example

>>> from rubix.core.ifu import get_doppler_shift_and_resampling
>>> doppler_shift_and_resampling = get_doppler_shift_and_resampling(config)

>>> rubixdata = doppler_shift_and_resampling(rubixdata)
>>> # Access the spectra of the stars, which is now doppler shifted and resampled to the telescope wavelength grid
>>> rubixdata.stars.spectra
rubix.core.ifu.get_resample_spectrum_pmap(target_wavelength) → Callable[source]#

Pmap the function that resamples the spectra of the stars to the telescope wavelength grid.

Parameters:

target_wavelength (jax.Array) – The telescope wavelength grid

Returns:

The function that resamples the spectra to the telescope wavelength grid.

rubix.core.ifu.get_resample_spectrum_vmap(target_wavelength) → Callable[source]#

The spectra of the stars are resampled to the telescope wavelength grid.

Parameters:

target_wavelength (jax.Array) – The telescope wavelength grid

Returns:

The function that resamples the spectra to the telescope wavelength grid.

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

The spectra of the stellar particles are scaled by the mass of the stars.

Parameters:

config (dict) – The configuration dictionary

Returns:

The function that scales the spectra by the mass of the stars.

Example

>>> from rubix.core.ifu import get_scale_spectrum_by_mass
>>> scale_spectrum_by_mass = get_scale_spectrum_by_mass(config)

>>> rubixdata = scale_spectrum_by_mass(rubixdata)
>>> # Access the spectra of the stars, which is now scaled by the stellar mass
>>> rubixdata.stars.spectra
rubix.core.ifu.get_velocities_doppler_shift_vmap(ssp_wave: Float[Array, '...'], velocity_direction: str) → Callable[source]#

The function doppler shifts the wavelength based on the velocity of the stars.

Parameters:

  • ssp_wave (jax.Array) – The wavelength of the SSP grid

  • velocity_direction (str) – The velocity component of the stars that is used to doppler shift the wavelength

Returns:

The function that doppler shifts the wavelength based on the velocity of the stars.

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"
>>> pipeline = RubixPipeline(config)
>>> output = pipeline.run()
>>> ssp_model = pipeline.ssp
>>> telescope = pipeline.telescope
gradient()[source]#

This function will calculate the gradient of the pipeline, but is yet not implemented.

run()[source]#

Runs the data processing pipeline.

Returns:

Output of the pipeline after processing the input data.

Return type:

dict

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) → Array[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