diff --git a/notebooks/rubix_pipeline_nihao.ipynb b/notebooks/rubix_pipeline_nihao.ipynb index 9455ab75..17462920 100644 --- a/notebooks/rubix_pipeline_nihao.ipynb +++ b/notebooks/rubix_pipeline_nihao.ipynb @@ -1,5 +1,18 @@ { "cells": [ + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "# NBVAL_SKIP\n", + "import os\n", + "os.environ['SPS_HOME'] = '/mnt/storage/annalena_data/sps_fsps'\n", + "#os.environ['SPS_HOME'] = '/home/annalena/sps_fsps'\n", + "#os.environ['SPS_HOME'] = '/Users/annalena/Documents/GitHub/fsps'" + ] + }, { "cell_type": "markdown", "metadata": {}, @@ -41,13 +54,13 @@ " \"save_data_path\": \"data\",\n", " },\n", " \"load_galaxy_args\": {\"reuse\": True},\n", - " \"subset\": {\"use_subset\": True, \"subset_size\": 1000},\n", + " \"subset\": {\"use_subset\": False, \"subset_size\": 1000},\n", " },\n", " \"simulation\": {\n", " \"name\": \"NIHAO\",\n", " \"args\": {\n", - " \"path\": \"/mnt/storage/_data/nihao/nihao_classic/g7.55e11/g7.55e11.01024\",\n", - " \"halo_path\": \"/mnt/storage/_data/nihao/nihao_classic/g7.55e11/g7.55e11.01024.z0.000.AHF_halos\",\n", + " \"path\": \"/mnt/storage/_data/nihao/nihao_classic/g8.26e11/g8.26e11.01024\",\n", + " \"halo_path\": \"/mnt/storage/_data/nihao/nihao_classic/g8.26e11/g8.26e11.01024.z0.000.AHF_halos\",\n", " #\"path\": \"/home/annalena/g7.55e11/snap_1024/output/7.55e11.01024\",\n", " #\"halo_path\": \"/home/annalena/g7.55e11/snap_1024/output/7.55e11.01024.z0.000.AHF_halos\",\n", " \"halo_id\": 0,\n", @@ -56,18 +69,25 @@ " \"output_path\": \"output\",\n", "\n", " \"telescope\": {\n", - " \"name\": \"MUSE\",\n", + " \"name\": \"MUSE_WFM\",\n", " \"psf\": {\"name\": \"gaussian\", \"size\": 5, \"sigma\": 0.6},\n", - " \"lsf\": {\"sigma\": 0.5},\n", - " \"noise\": {\"signal_to_noise\": 1, \"noise_distribution\": \"normal\"},\n", + " \"lsf\": {\"sigma\": 1.2},\n", + " \"noise\": {\"signal_to_noise\": 100, \"noise_distribution\": \"normal\"},\n", " },\n", " \"cosmology\": {\"name\": \"PLANCK15\"},\n", " \"galaxy\": {\n", - " \"dist_z\": 0.2,\n", + " \"dist_z\": 0.01,\n", " \"rotation\": {\"type\": \"edge-on\"},\n", " },\n", " \"ssp\": {\n", - " \"template\": {\"name\": \"BruzualCharlot2003\"},\n", + " \"template\": {\"name\": \"FSPS\"},\n", + " \"dust\": {\n", + " \"extinction_model\": \"Cardelli89\",\n", + " \"dust_to_gas_ratio\": 0.01,\n", + " \"dust_to_metals_ratio\": 0.4,\n", + " \"dust_grain_density\": 3.5,\n", + " \"Rv\": 3.1,\n", + " },\n", " },\n", "}" ] @@ -119,7 +139,7 @@ "wave = pipe.telescope.wave_seq\n", "spectra = rubixdata.stars.datacube\n", "\n", - "plt.plot(wave, spectra[12, 12, :])\n", + "plt.plot(wave, spectra[120, 120, :])\n", "plt.title(\"Spectrum of Spaxel [12, 12]\")\n", "plt.xlabel(\"Wavelength [Å]\")\n", "plt.ylabel(\"Flux\")\n", @@ -170,6 +190,114 @@ "stellar_age_histogram('./output/rubix_galaxy.h5')" ] }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "### Mean line of sight velocity" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "# NBVAL_SKIP\n", + "import numpy as np\n", + "import matplotlib.pyplot as plt\n", + "\n", + "# Assuming your data arrays are defined as follows:\n", + "pixel_assignment = np.asarray(np.squeeze(rubixdata.stars.pixel_assignment))\n", + "velocities = np.asarray(rubixdata.stars.velocity[0, :, 2])\n", + "\n", + "# Compute the sum of velocities and count per pixel using np.bincount\n", + "sum_velocity = np.bincount(pixel_assignment, weights=velocities)\n", + "counts = np.bincount(pixel_assignment)\n", + "\n", + "# Calculate mean velocity; note: division by zero is avoided if every pixel has at least one star.\n", + "mean_velocity = sum_velocity / counts\n", + "\n", + "\n", + "# If you know the pixel grid dimensions (for example, a square grid)\n", + "n_pixels = len(mean_velocity)\n", + "grid_size = int(np.sqrt(n_pixels))\n", + "if grid_size * grid_size != n_pixels:\n", + " raise ValueError(\"The total number of pixels is not a perfect square; please specify the grid shape explicitly.\")\n", + "\n", + "# Reshape the mean_velocity into a 2D array for imshow\n", + "velocity_map = mean_velocity.reshape((grid_size, grid_size))\n", + "print(velocity_map[12,12])\n", + "\n", + "print(velocity_map[17,12]-velocity_map[7,12])\n", + "# Plot the result\n", + "plt.figure(figsize=(6, 5))\n", + "plt.imshow(velocity_map, origin='lower', interpolation='nearest', cmap='seismic')\n", + "plt.colorbar(label='Mean Velocity')\n", + "plt.title('Mean Velocity per Pixel')\n", + "plt.xlabel('X pixel index')\n", + "plt.ylabel('Y pixel index')\n", + "#storepath = f\"output/datacube_NIHAO{config['data']['load_galaxy_args']['id']}_{config['pipeline']['name']}_velocity.png\"\n", + "#plt.savefig(storepath)\n", + "plt.show()\n" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "### Mean stellar age" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "\n", + "\n", + "# NBVAL_SKIP\n", + "import numpy as np\n", + "import matplotlib.pyplot as plt\n", + "\n", + "# Assuming your data arrays are defined as follows:\n", + "pixel_assignment = np.asarray(np.squeeze(rubixdata.stars.pixel_assignment))\n", + "ages = np.asarray(rubixdata.stars.age[0, :])\n", + "\n", + "# Compute the sum of velocities and count per pixel using np.bincount\n", + "sum_ages = np.bincount(pixel_assignment, weights=ages)\n", + "counts = np.bincount(pixel_assignment)\n", + "\n", + "# Calculate mean velocity; note: division by zero is avoided if every pixel has at least one star.\n", + "mean_age = sum_ages / counts\n", + "\n", + "\n", + "# If you know the pixel grid dimensions (for example, a square grid)\n", + "n_pixels = len(mean_age)\n", + "grid_size = int(np.sqrt(n_pixels))\n", + "if grid_size * grid_size != n_pixels:\n", + " raise ValueError(\"The total number of pixels is not a perfect square; please specify the grid shape explicitly.\")\n", + "\n", + "# Reshape the mean_velocity into a 2D array for imshow\n", + "age_map = mean_age.reshape((grid_size, grid_size))\n", + "print(age_map[12,12])\n", + "\n", + "# Plot the result\n", + "plt.figure(figsize=(6, 5))\n", + "plt.imshow(age_map, origin='lower', interpolation='nearest', cmap='inferno')\n", + "plt.colorbar(label='Mean Age')\n", + "plt.title('Mean Age per Pixel')\n", + "plt.xlabel('X pixel index')\n", + "plt.ylabel('Y pixel index')\n", + "#storepath = f\"./output/datacube_NIHAO{config['data']['load_galaxy_args']['id']}_{config[\"telescope\"][\"name\"]}_{config['pipeline']['name']}_age.png\"\n", + "#plt.savefig(storepath)\n", + "plt.show()\n", + "\n", + "\n", + "\n" + ] + }, { "cell_type": "markdown", "metadata": {}, @@ -182,7 +310,7 @@ ], "metadata": { "kernelspec": { - "display_name": "rubix", + "display_name": "Python 3", "language": "python", "name": "python3" }, @@ -196,7 +324,7 @@ "name": "python", "nbconvert_exporter": "python", "pygments_lexer": "ipython3", - "version": "3.10.13" + "version": "3.13.2" } }, "nbformat": 4, diff --git a/rubix/config/pynbody_config.yml b/rubix/config/pynbody_config.yml index 2d9ff362..d25f0459 100644 --- a/rubix/config/pynbody_config.yml +++ b/rubix/config/pynbody_config.yml @@ -9,7 +9,10 @@ fields: gas: density: "rho" temperature: "temp" - metallicity: "metals" + metals: "metals" + #OxMassFrac: "OxMassFrac" + #HI: "HI" + metallicity: metals coords: "pos" velocity: "vel" mass: "mass" @@ -28,6 +31,9 @@ units: gas: density: "Msun/kpc^3" temperature: "K" + metals: "dimensionless" + #OxMassFrac: "dimensionless" + #HI: "dimensionless" metallicity: "Zsun" coords: "kpc" velocity: "km/s" diff --git a/rubix/core/fits.py b/rubix/core/fits.py index 7f2ff875..a766122d 100644 --- a/rubix/core/fits.py +++ b/rubix/core/fits.py @@ -41,7 +41,15 @@ def store_fits(config, data, filepath): hdr["SIMPLE"] = "T /conforms to FITS standard" hdr["PIPELINE"] = config["pipeline"]["name"] hdr["DIST_z"] = config["galaxy"]["dist_z"] - hdr["ROTATION"] = config["galaxy"]["rotation"]["type"] + if ( + config["galaxy"]["rotation"]["type"] == "face-on" + or config["galaxy"]["rotation"]["type"] == "edge-on" + ): + hdr["ROTATION"] = config["galaxy"]["rotation"]["type"] + else: + hdr["ROT_a"] = config["galaxy"]["rotation"]["alpha"] + hdr["ROT_b"] = config["galaxy"]["rotation"]["beta"] + hdr["ROT_c"] = config["galaxy"]["rotation"]["gamma"] hdr["SIM"] = config["simulation"]["name"] # For Illustris and NIHAO @@ -66,7 +74,7 @@ def store_fits(config, data, filepath): hdr1 = fits.Header() hdr1["EXTNAME"] = "DATA" hdr1["OBJECT"] = object_name - hdr1["BUNIT"] = "erg/(s*cm^2*A)" # flux unit per Angstrom + hdr1["BUNIT"] = "10**-20 erg/(s*cm^2*A)" # flux unit per Angstrom hdr1["CRPIX1"] = (datacube.shape[0] - 1) / 2 hdr1["CRPIX2"] = (datacube.shape[1] - 1) / 2 hdr1["CD1_1"] = telescope.spatial_res / 3600 # convert arcsec to deg @@ -92,7 +100,7 @@ def store_fits(config, data, filepath): output_filename = ( f"{filepath}{config['simulation']['name']}_id{galaxy_id}_snap{snapshot}_" - f"{parttype}_subset{config['data']['subset']['use_subset']}.fits" + f'{config["telescope"]["name"]}_{config["pipeline"]["name"]}.fits' ) os.makedirs(os.path.dirname(output_filename), exist_ok=True) diff --git a/rubix/core/ifu.py b/rubix/core/ifu.py index c41a3ce7..2388f76d 100644 --- a/rubix/core/ifu.py +++ b/rubix/core/ifu.py @@ -93,7 +93,7 @@ def calculate_spectra(rubixdata: RubixData) -> RubixData: spectra = jnp.concatenate([spectra1, spectra2, spectra3], axis=0) """ # Define the chunk size (number of particles per chunk) - chunk_size = 250000 + chunk_size = 100000 total_length = metallicity[0].shape[ 0 ] # assuming metallicity[0] is your 1D array of particles diff --git a/rubix/core/rotation.py b/rubix/core/rotation.py index 8b1df89a..05c8d3df 100644 --- a/rubix/core/rotation.py +++ b/rubix/core/rotation.py @@ -76,9 +76,11 @@ def get_galaxy_rotation(config: dict): def rotate_galaxy(rubixdata: RubixData) -> RubixData: logger.info(f"Rotating galaxy with alpha={alpha}, beta={beta}, gamma={gamma}") + """ for particle_type in ["stars", "gas"]: if particle_type in config["data"]["args"]["particle_type"]: # Get the component (either stars or gas) + logger.info(f"Rotating {particle_type}") component = getattr(rubixdata, particle_type) # Get the inputs @@ -99,7 +101,8 @@ def rotate_galaxy(rubixdata: RubixData) -> RubixData: coords, velocities = rotate_galaxy_core( positions=coords, velocities=velocities, - masses=masses, + positions_stars=rubixdata.stars.coords, + masses_stars=rubixdata.stars.mass, halfmass_radius=halfmass_radius, alpha=alpha, beta=beta, @@ -113,5 +116,64 @@ def rotate_galaxy(rubixdata: RubixData) -> RubixData: setattr(component, "velocity", velocities) return rubixdata + """ + logger.info("Rotating galaxy for simulation: " + config["simulation"]["name"]) + # Rotate gas + if "gas" in config["data"]["args"]["particle_type"]: + logger.info("Rotating gas") + + # Rotate the gas component + new_coords_gas, new_velocities_gas = rotate_galaxy_core( + positions=rubixdata.gas.coords, + velocities=rubixdata.gas.velocity, + positions_stars=rubixdata.stars.coords, + masses_stars=rubixdata.stars.mass, + halfmass_radius=rubixdata.galaxy.halfmassrad_stars, + alpha=alpha, + beta=beta, + gamma=gamma, + key=config["simulation"]["name"], + ) + + setattr(rubixdata.gas, "coords", new_coords_gas) + setattr(rubixdata.gas, "velocity", new_velocities_gas) + + # Rotate the stellar component + new_coords_stars, new_velocities_stars = rotate_galaxy_core( + positions=rubixdata.stars.coords, + velocities=rubixdata.stars.velocity, + positions_stars=rubixdata.stars.coords, + masses_stars=rubixdata.stars.mass, + halfmass_radius=rubixdata.galaxy.halfmassrad_stars, + alpha=alpha, + beta=beta, + gamma=gamma, + key=config["simulation"]["name"], + ) + + setattr(rubixdata.stars, "coords", new_coords_stars) + setattr(rubixdata.stars, "velocity", new_velocities_stars) + + else: + logger.warning( + "Gas not found in particle_type, only rotating stellar component." + ) + # Rotate the stellar component + new_coords_stars, new_velocities_stars = rotate_galaxy_core( + positions=rubixdata.stars.coords, + velocities=rubixdata.stars.velocity, + positions_stars=rubixdata.stars.coords, + masses_stars=rubixdata.stars.mass, + halfmass_radius=rubixdata.galaxy.halfmassrad_stars, + alpha=alpha, + beta=beta, + gamma=gamma, + key=config["simulation"]["name"], + ) + + setattr(rubixdata.stars, "coords", new_coords_stars) + setattr(rubixdata.stars, "velocity", new_velocities_stars) + + return rubixdata return rotate_galaxy diff --git a/rubix/galaxy/alignment.py b/rubix/galaxy/alignment.py index 1398384a..09b41538 100644 --- a/rubix/galaxy/alignment.py +++ b/rubix/galaxy/alignment.py @@ -233,11 +233,13 @@ def apply_rotation( def rotate_galaxy( positions: Float[Array, "* 3"], velocities: Float[Array, "* 3"], - masses: Float[Array, "..."], + positions_stars: Float[Array, "..."], + masses_stars: Float[Array, "..."], halfmass_radius: Union[Float[Array, "..."], float], alpha: float, beta: float, gamma: float, + key: str, ) -> Tuple[Float[Array, "* 3"], Float[Array, "* 3"]]: """ Orientate the galaxy by applying a rotation matrix to the positions of the particles. @@ -254,12 +256,25 @@ def rotate_galaxy( Returns: The rotated positions and velocities as a jnp.ndarray. """ - - I = moment_of_inertia_tensor(positions, masses, halfmass_radius) - R = rotation_matrix_from_inertia_tensor(I) - pos_rot = apply_init_rotation(positions, R) - vel_rot = apply_init_rotation(velocities, R) - pos_final = apply_rotation(pos_rot, alpha, beta, gamma) - vel_final = apply_rotation(vel_rot, alpha, beta, gamma) + # we have to distinguis between IllustrisTNG and NIHAO. + # The nihao galaxies are already oriented face-on in the pynbody input handler. + # The IllustrisTNG galaxies are not oriented face-on, so we have to calculate the moment of inertia tensor + # and apply the rotation matrix to the positions and velocities. + # After that the simulations can be treated in the same way. + # Then the user specific rotation is applied to the positions and velocities. + if key == "IllustrisTNG": + I = moment_of_inertia_tensor(positions_stars, masses_stars, halfmass_radius) + R = rotation_matrix_from_inertia_tensor(I) + pos_rot = apply_init_rotation(positions, R) + vel_rot = apply_init_rotation(velocities, R) + pos_final = apply_rotation(pos_rot, alpha, beta, gamma) + vel_final = apply_rotation(vel_rot, alpha, beta, gamma) + elif key == "NIHAO": + pos_final = apply_rotation(positions, alpha, beta, gamma) + vel_final = apply_rotation(velocities, alpha, beta, gamma) + else: + raise ValueError( + f"Unknown key: {key} for the rotation. Supported keys are 'IllustrisTNG' and 'NIHAO'." + ) return pos_final, vel_final diff --git a/rubix/galaxy/input_handler/pynbody.py b/rubix/galaxy/input_handler/pynbody.py index 5a5ebf57..1fc1f29e 100644 --- a/rubix/galaxy/input_handler/pynbody.py +++ b/rubix/galaxy/input_handler/pynbody.py @@ -6,6 +6,7 @@ import pynbody import yaml +from rubix.cosmology import PLANCK15 as rubix_cosmo from rubix.units import Zsun from rubix.utils import SFTtoAge @@ -88,6 +89,36 @@ def load_data(self): getattr(self.sim, cls), fields[cls], units[cls], cls ) + # Combine HI and OxMassFrac into a two-column metals field for gas + hi_data = self.load_particle_data( + getattr(self.sim, "gas"), + {"HI": "HI"}, + {"HI": u.dimensionless_unscaled}, + "gas", + ) + ox_data = self.load_particle_data( + getattr(self.sim, "gas"), + {"OxMassFrac": "OxMassFrac"}, + {"OxMassFrac": u.dimensionless_unscaled}, + "gas", + ) + # fe_data = self.load_particle_data(getattr(self.sim, "gas"), {"FeMassFrac": "FeMassFrac"}, {"FeMassFrac": u.dimensionless_unscaled}, "gas") + # self.data["gas"]["metals"] = np.column_stack((hi_data["HI"], ox_data["OxMassFrac"])) + # Create a metals array with 10 columns, filled with zeros initially + n_particles = hi_data["HI"].shape[0] + metals = np.zeros((n_particles, 10), dtype=hi_data["HI"].dtype) + + # Place HI values at column 0 and OxMassFrac (O) at column 4 (that it is storred in the same way as IllustrisTNG) + metals[:, 0] = hi_data["HI"] + metals[:, 4] = ox_data["OxMassFrac"] + + self.data["gas"]["metals"] = metals + self.logger.info("Metals assigned to gas particles.") + self.logger.info("Metals shape is: %s", self.data["gas"]["metals"].shape) + + age_at_z0 = rubix_cosmo.age_at_z0() + self.data["stars"]["age"] = age_at_z0 * u.Gyr - self.data["stars"]["age"] + self.logger.info( f"Simulation snapshot and halo data loaded successfully for classes: {load_classes}." ) diff --git a/rubix/telescope/telescopes.yaml b/rubix/telescope/telescopes.yaml index 73300afe..1f191807 100644 --- a/rubix/telescope/telescopes.yaml +++ b/rubix/telescope/telescopes.yaml @@ -9,9 +9,19 @@ MUSE: aperture_type: "square" pixel_type: "square" -MUSE_NFM: - fov: 7.5 - spatial_res: 0.05 +MUSE_WFM: + fov: 60.0 + spatial_res: 0.2 + wave_range: [4700.15, 9351.4] + wave_res: 1.25 + lsf_fwhm: 2.51 + signal_to_noise: null + aperture_type: "square" + pixel_type: "square" + +MUSE_ultraWFM: + fov: 180.0 + spatial_res: 0.2 wave_range: [4700.15, 9351.4] wave_res: 1.25 lsf_fwhm: 2.51 diff --git a/tests/test_galaxy_alignment.py b/tests/test_galaxy_alignment.py index 173dd149..74521d3b 100644 --- a/tests/test_galaxy_alignment.py +++ b/tests/test_galaxy_alignment.py @@ -187,7 +187,15 @@ def test_rotate_galaxy(): gamma = 0.0 rotated_positions, rotated_velocities = rotate_galaxy( - positions, velocities, masses, halfmass_radius, alpha, beta, gamma + positions, + velocities, + positions, + masses, + halfmass_radius, + alpha, + beta, + gamma, + "IllustrisTNG", ) assert rotated_positions.shape == positions.shape