RUBIX pipeline#
RUBIX is designed as a linear pipeline, where the individual functions are called and constructed as a pipeline. This allows as to execude the whole data transformation from a cosmological hydrodynamical simulation of a galaxy to an IFU cube in two lines of code. This notebook shows, how to execute the pipeline. To see, how the pipeline is execuded in small individual steps per individual function, we refer to the notebook rubix_pipeline_stepwise.ipynb.
How to use the Pipeline#
Define a
configSetup the
pipeline yamlRun the RUBIX pipeline
Do science with the mock-data
Step 1: Config#
The config contains all the information needed to run the pipeline. Those are run specfic configurations. Currently we just support Illustris as simulation, but extensions to other simulations (e.g. NIHAO) are planned.
For the config you can choose the following options:
pipeline: you specify the name of the pipeline that is stored in the yaml file in rubix/config/pipeline_config.ymllogger: RUBIX has implemented a logger to report the user, what is happening during the pipeline execution and give warningsdata - args - particle_type: load only stars particle (“particle_type”: [“stars”]) or only gas particle (“particle_type”: [“gas”]) or both (“particle_type”: [“stars”,”gas”])data - args - simulation: choose the Illustris simulation (e.g. “simulation”: “TNG50-1”)data - args - snapshot: which time step of the simulation (99 for present day)data - args - save_data_path: set the path to save the downloaded Illustris datadata - load_galaxy_args - id: define, which Illustris galaxy is downloadeddata - load_galaxy_args - reuse: if True, if in th esave_data_path directory a file for this galaxy id already exists, the downloading is skipped and the preexisting file is useddata - subset: only a defined number of stars/gas particles is used and stored for the pipeline. This may be helpful for quick testingsimulation - name: currently only IllustrisTNG is supportedsimulation - args - path: where the data is stored and how the file will be namedoutput_path: where the hdf5 file is stored, which is then the input to the RUBIX pipelinetelescope - name: define the telescope instrument that is observing the simulation. Some telescopes are predefined, e.g. MUSE. If your instrument does not exist predefined, you can easily define your instrument in rubix/telescope/telescopes.yamltelescope - psf: define the point spread function that is applied to the mock datatelescope - lsf: define the line spread function that is applied to the mock datatelescope - noise: define the noise that is applied to the mock datacosmology: specify the cosmology you want to use, standard for RUBIX is “PLANCK15”galaxy - dist_z: specify at which redshift the mock-galaxy is observedgalaxy - rotation: specify the orientation of the galaxy. You can set the types edge-on or face-on or specify the angles alpha, beta and gamma as rotations around x-, y- and z-axisssp - template: specify the simple stellar population lookup template to get the stellar spectrum for each stars particle. In RUBIX frequently “BruzualCharlot2003” is used.
# NBVAL_SKIP
from jax import config
import os
import jax
print(jax.devices())
[CpuDevice(id=0)]
# NBVAL_SKIP
import os
os.environ['SPS_HOME'] = '/home/annalena/sps_fsps'
#NBVAL_SKIP
import matplotlib.pyplot as plt
from rubix.core.pipeline import RubixPipeline
import os
galaxy_id = "g8.13e11"
config_NIHAO = {
"pipeline":{"name": "calc_ifu"},
"logger": {
"log_level": "DEBUG",
"log_file_path": None,
"format": "%(asctime)s - %(name)s - %(levelname)s - %(message)s",
},
"data": {
"name": "NihaoHandler",
"args": {
"particle_type": ["stars"],
"save_data_path": "data",
"snapshot": "1024",
},
"load_galaxy_args": {"reuse": True, "id": galaxy_id},
"subset": {"use_subset": False, "subset_size": 200000},
},
"simulation": {
"name": "NIHAO",
"args": {
"path": f'/home/_data/nihao/nihao_classic/{galaxy_id}/{galaxy_id}.01024',
"halo_path": f'/home/_data/nihao/nihao_classic/{galaxy_id}/{galaxy_id}.01024.z0.000.AHF_halos',
"halo_id": 0,
},
},
"output_path": "output",
"telescope":
{"name": "MUSE",
"psf": {"name": "gaussian", "size": 5, "sigma": 0.6},
"lsf": {"sigma": 0.5},
"noise": {"signal_to_noise": 100,"noise_distribution": "normal"},},
"cosmology":
{"name": "PLANCK15"},
"galaxy":
{"dist_z": 0.1,
"rotation": {"type": "edge-on"},
},
"ssp": {
"template": {
"name": "Mastar_CB19_SLOG_1_5"
},
"dust": {
"extinction_model": "Cardelli89",
"dust_to_gas_ratio": 0.01,
"dust_to_metals_ratio": 0.4,
"dust_grain_density": 3.5,
"Rv": 3.1,
},
},
}
2025-11-11 10:23:01,492 - rubix - INFO -
___ __ _____ _____ __
/ _ \/ / / / _ )/ _/ |/_/
/ , _/ /_/ / _ |/ /_> <
/_/|_|\____/____/___/_/|_|
2025-11-11 10:23:01,494 - rubix - INFO - Rubix version: 0.0.post626+g42b4b7505.d20251110
2025-11-11 10:23:01,494 - rubix - INFO - JAX version: 0.7.2
2025-11-11 10:23:01,495 - rubix - INFO - Running on [CpuDevice(id=0)] devices
2025-11-11 10:23:01,496 - rubix - WARNING - python-fsps is not installed. Please install it to use this function. Install using pip install fsps and check the installation page: https://dfm.io/python-fsps/current/installation/ for more details. Especially, make sure to set all necessary environment variables.
# NBVAL_SKIP
config_TNG = {
"pipeline":{"name": "calc_ifu"},
"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"],
"simulation": "TNG50-1",
"snapshot": 99,
"save_data_path": "data",
},
"load_galaxy_args": {
"id": 12,
"reuse": True,
},
"subset": {
"use_subset": True,
"subset_size": 2000,
},
},
"simulation": {
"name": "IllustrisTNG",
"args": {
"path": "data/galaxy-id-12.hdf5",
},
},
"output_path": "output",
"telescope":
{"name": "MUSE",
"psf": {"name": "gaussian", "size": 5, "sigma": 0.6},
"lsf": {"sigma": 0.5},
"noise": {"signal_to_noise": 100,"noise_distribution": "normal"},},
"cosmology":
{"name": "PLANCK15"},
"galaxy":
{"dist_z": 0.1,
"rotation": {"type": "edge-on"},
},
"ssp": {
"template": {
"name": "Mastar_CB19_SLOG_1_5"
},
"dust": {
"extinction_model": "Cardelli89",
"dust_to_gas_ratio": 0.01,
"dust_to_metals_ratio": 0.4,
"dust_grain_density": 3.5,
"Rv": 3.1,
},
},
}
Step 2: Pipeline yaml#
To run the RUBIX pipeline, you need a yaml file (stored in rubix/config/pipeline_config.yml) that defines which functions are used during the execution of the pipeline. This shows the example pipeline yaml to compute a stellar IFU cube.
calc_ifu:
Transformers:
rotate_galaxy:
name: rotate_galaxy
depends_on: null
args: []
kwargs:
type: "face-on"
filter_particles:
name: filter_particles
depends_on: rotate_galaxy
args: []
kwargs: {}
spaxel_assignment:
name: spaxel_assignment
depends_on: filter_particles
args: []
kwargs: {}
reshape_data:
name: reshape_data
depends_on: spaxel_assignment
args: []
kwargs: {}
calculate_spectra:
name: calculate_spectra
depends_on: reshape_data
args: []
kwargs: {}
scale_spectrum_by_mass:
name: scale_spectrum_by_mass
depends_on: calculate_spectra
args: []
kwargs: {}
doppler_shift_and_resampling:
name: doppler_shift_and_resampling
depends_on: scale_spectrum_by_mass
args: []
kwargs: {}
calculate_datacube:
name: calculate_datacube
depends_on: doppler_shift_and_resampling
args: []
kwargs: {}
convolve_psf:
name: convolve_psf
depends_on: calculate_datacube
args: []
kwargs: {}
convolve_lsf:
name: convolve_lsf
depends_on: convolve_psf
args: []
kwargs: {}
apply_noise:
name: apply_noise
depends_on: convolve_lsf
args: []
kwargs: {}
Ther is one thing you have to know about the naming of the functions in this yaml: To use the functions inside the pipeline, the functions have to be called exactly the same as they are returned from the core module function!
Step 3: Run the pipeline#
After defining the config and the pipeline_config you can simply run the whole pipeline by these two lines of code.
#NBVAL_SKIP
pipe = RubixPipeline(config_TNG)
---------------------------------------------------------------------------
AttributeError Traceback (most recent call last)
Cell In[6], line 2
1 #NBVAL_SKIP
----> 2 pipe = RubixPipeline(config_TNG)
File ~/.conda/envs/rubix/lib/python3.12/site-packages/rubix/core/pipeline.py:87, in RubixPipeline.__init__(self, user_config)
85 self.user_config = get_config(user_config)
86 self.pipeline_config = get_pipeline_config(self.user_config["pipeline"]["name"])
---> 87 self.logger = get_logger(self.user_config["logger"])
88 self.ssp = get_ssp(self.user_config)
89 self.telescope = get_telescope(self.user_config)
File ~/.conda/envs/rubix/lib/python3.12/site-packages/rubix/logger.py:21, in get_logger(config)
18 else:
19 first_time = False
---> 21 logger.setLevel(getattr(logging, config["log_level"].upper(), "INFO"))
23 # Clear existing handlers
24 for handler in logger.handlers[:]:
AttributeError: 'NoneType' object has no attribute 'upper'
#NBVAL_SKIP
devices = jax.devices()
inputdata = pipe.prepare_data()
rubixdata = pipe.run_sharded(inputdata, devices)
2025-11-11 10:22:23,162 - rubix - INFO - Getting rubix data...
2025-11-11 10:22:23,163 - rubix - INFO - Rubix galaxy file already exists, skipping conversion
/home/annalena/.conda/envs/rubix/lib/python3.12/site-packages/jax/_src/numpy/scalar_types.py:50: UserWarning: Explicitly requested dtype float64 requested in asarray is not available, and will be truncated to dtype float32. To enable more dtypes, set the jax_enable_x64 configuration option or the JAX_ENABLE_X64 shell environment variable. See https://github.com/jax-ml/jax#current-gotchas for more.
return asarray(x, dtype=self.dtype)
/home/annalena/.conda/envs/rubix/lib/python3.12/site-packages/rubix/core/data.py:491: UserWarning: Explicitly requested dtype <class 'jax.numpy.float64'> requested in array is not available, and will be truncated to dtype float32. To enable more dtypes, set the jax_enable_x64 configuration option or the JAX_ENABLE_X64 shell environment variable. See https://github.com/jax-ml/jax#current-gotchas for more.
rubixdata.galaxy.center = jnp.array(data["subhalo_center"], dtype=jnp.float64)
2025-11-11 10:22:23,246 - rubix - INFO - Centering stars particles
2025-11-11 10:22:24,250 - rubix - WARNING - The Subset value is set in config. Using only subset of size 2000 for stars
2025-11-11 10:22:24,254 - rubix - INFO - Data loaded with 2000 star particles and 0 gas particles.
2025-11-11 10:22:24,255 - rubix - INFO - Data preparation completed in 1.09 seconds.
2025-11-11 10:22:24,256 - rubix - INFO - Setting up the pipeline...
2025-11-11 10:22:24,257 - rubix - DEBUG - Pipeline Configuration: {'Transformers': {'rotate_galaxy': {'name': 'rotate_galaxy', 'depends_on': None, 'args': [], 'kwargs': {}}, 'filter_particles': {'name': 'filter_particles', 'depends_on': 'rotate_galaxy', 'args': [], 'kwargs': {}}, 'spaxel_assignment': {'name': 'spaxel_assignment', 'depends_on': 'filter_particles', 'args': [], 'kwargs': {}}, 'calculate_datacube_particlewise': {'name': 'calculate_datacube_particlewise', 'depends_on': 'spaxel_assignment', 'args': [], 'kwargs': {}}, 'convolve_psf': {'name': 'convolve_psf', 'depends_on': 'calculate_datacube_particlewise', 'args': [], 'kwargs': {}}, 'convolve_lsf': {'name': 'convolve_lsf', 'depends_on': 'convolve_psf', 'args': [], 'kwargs': {}}, 'apply_noise': {'name': 'apply_noise', 'depends_on': 'convolve_lsf', 'args': [], 'kwargs': {}}}}
2025-11-11 10:22:24,258 - rubix - DEBUG - Rotation Type found: edge-on
2025-11-11 10:22:24,261 - rubix - INFO - Calculating spatial bin edges...
/home/annalena/.conda/envs/rubix/lib/python3.12/site-packages/rubix/telescope/factory.py:26: UserWarning: No telescope config provided, using default stored in /home/annalena/.conda/envs/rubix/lib/python3.12/site-packages/rubix/telescope/telescopes.yaml
warnings.warn(
2025-11-11 10:22:24,285 - rubix - INFO - Getting cosmology...
2025-11-11 10:22:24,497 - rubix - INFO - Calculating spatial bin edges...
2025-11-11 10:22:24,509 - rubix - INFO - Getting cosmology...
2025-11-11 10:22:24,521 - rubix - INFO - Getting cosmology...
2025-11-11 10:22:25,075 - rubix - DEBUG - Method not defined, using default method: cubic
/home/annalena/.conda/envs/rubix/lib/python3.12/site-packages/rubix/telescope/factory.py:26: UserWarning: No telescope config provided, using default stored in /home/annalena/.conda/envs/rubix/lib/python3.12/site-packages/rubix/telescope/telescopes.yaml
warnings.warn(
2025-11-11 10:22:26,192 - rubix - DEBUG - Method not defined, using default method: cubic
/home/annalena/.conda/envs/rubix/lib/python3.12/site-packages/rubix/telescope/factory.py:26: UserWarning: No telescope config provided, using default stored in /home/annalena/.conda/envs/rubix/lib/python3.12/site-packages/rubix/telescope/telescopes.yaml
warnings.warn(
2025-11-11 10:22:26,935 - rubix - INFO - Assembling the pipeline...
2025-11-11 10:22:26,936 - rubix - INFO - Compiling the expressions...
2025-11-11 10:22:26,936 - rubix - INFO - Number of devices: 1
2025-11-11 10:22:27,024 - rubix - INFO - Rotating galaxy with alpha=90.0, beta=0.0, gamma=0.0
2025-11-11 10:22:27,024 - rubix - INFO - Rotating galaxy for simulation: IllustrisTNG
2025-11-11 10:22:27,025 - rubix - WARNING - Gas not found in particle_type, only rotating stellar component.
2025-11-11 10:22:27,143 - rubix - INFO - Filtering particles outside the aperture...
2025-11-11 10:22:27,150 - rubix - INFO - Assigning particles to spaxels...
2025-11-11 10:22:27,184 - rubix - INFO - Calculating Data Cube (combined per‐particle)…
2025-11-11 10:22:27,479 - rubix - DEBUG - Datacube shape: (25, 25, 3721)
2025-11-11 10:22:27,480 - rubix - INFO - Convolving with PSF...
2025-11-11 10:22:27,484 - rubix - INFO - Convolving with LSF...
2025-11-11 10:22:27,490 - rubix - INFO - Applying noise to datacube with signal to noise ratio: 100 and noise distribution: normal
2025-11-11 10:22:29,524 - rubix - INFO - Total time for sharded pipeline run: 5.27 seconds.
#NBVAL_SKIP
#inputdata = pipe.prepare_data()
#shard_rubixdata = pipe.run_sharded_chunked(inputdata)
Step 4: Mock-data#
Now we have our final datacube and can use the mock-data to do science. Here we have a quick look in the optical wavelengthrange of the mock-datacube and show the spectra of a central spaxel and a spatial image.
#NBVAL_SKIP
import jax.numpy as jnp
wave = pipe.telescope.wave_seq
# get the indices of the visible wavelengths of 4000-8000 Angstroms
visible_indices = jnp.where((wave >= 4000) & (wave <= 8000))
This is how you can access the spectrum of an individual spaxel, the wavelength can be accessed via pipe.wave_seq
#NBVAL_SKIP
wave = pipe.telescope.wave_seq
#spectra = rubixdata#.stars.datacube # Spectra of all stars
spectra_sharded = rubixdata # Spectra of all stars
#print(spectra.shape)
plt.figure(figsize=(10, 5))
plt.title("Rubix Sharded")
plt.xlabel("Wavelength [Angstrom]")
plt.ylabel("Flux [erg/s/cm^2/Angstrom]")
plt.plot(wave, spectra_sharded[12,12,:])
plt.plot(wave, spectra_sharded[8,12,:])
plt.show()
Plot a spacial image of the data cube
#NBVAL_SKIP
# get the spectra of the visible wavelengths from the ifu cube
#visible_spectra = rubixdata.stars.datacube[ :, :, visible_indices[0]]
#visible_spectra = rubixdata[ :, :, visible_indices[0]]
sharded_visible_spectra = rubixdata[ :, :, visible_indices[0]]
#visible_spectra.shape
#image = jnp.sum(visible_spectra, axis=2)
sharded_image = jnp.sum(sharded_visible_spectra, axis=2)
# Plot side by side
fig, axes = plt.subplots(1, 1, figsize=(12, 5))
# Sharded IFU datacube image
im1 = axes.imshow(sharded_image, origin="lower", cmap="inferno")
axes.set_title("Sharded IFU Datacube")
fig.colorbar(im1, ax=axes)
plt.tight_layout()
plt.show()
DONE!#
Congratulations, you have sucessfully run the RUBIX pipeline to create your own mock-observed IFU datacube! Now enjoy playing around with the RUBIX pipeline and enjoy doing amazing science with RUBIX :)