From 46bfa5e26c6519eef428ee173af5c631cb9dc197 Mon Sep 17 00:00:00 2001 From: anschaible Date: Tue, 1 Apr 2025 15:27:23 +0200 Subject: [PATCH 01/11] add metals in pynbody handler that are needed for dust computation, hi,ox are stored in the same positions in the array as it is done for lllustris --- notebooks/nohup.out | 204 ++++++++++++++++++++++++++ notebooks/rubix_pipeline_nihao.ipynb | 120 ++++++++++++++- rubix/config/pynbody_config.yml | 8 +- rubix/galaxy/input_handler/pynbody.py | 41 +++++- 4 files changed, 363 insertions(+), 10 deletions(-) create mode 100644 notebooks/nohup.out diff --git a/notebooks/nohup.out b/notebooks/nohup.out new file mode 100644 index 00000000..b7dea1b2 --- /dev/null +++ b/notebooks/nohup.out @@ -0,0 +1,204 @@ +Processing galaxy_id: g5.02e11 +Running Papermill for pipeline: calc_ifu with galaxy_id: g5.02e11 at telescope: MUSE_ultraWFM +Input Notebook: fits_file_nihao_script.ipynb +Output Notebook: fits_file_nihao_script_executed.ipynb +Passed unknown parameter: galaxy_id +Passed unknown parameter: pipeline_name +Passed unknown parameter: telescope_name +Input notebook does not contain a cell with tag 'parameters' + Executing: 0%| | 0/17 [00:00 1s: + + %gather.423 = f32[10,107,5994]{2,1,0} gather(f32[11,107,5994]{2,1,0} %constant.7504, s32[10,1]{1,0} %constant.7532), offset_dims={1,2}, collapsed_slice_dims={0}, start_index_map={0}, index_vector_dim=1, slice_sizes={1,107,5994}, metadata={op_name="jit()/jit(main)/jit(interp2d)/jit(fun)/jit(_take)/gather" source_file="/home/annalena/rubix/rubix/core/ifu.py" source_line=107} + +This isn't necessarily a bug; constant-folding is inherently a trade-off between compilation time and speed at runtime. XLA has some guards that attempt to keep constant folding from taking too long, but fundamentally you'll always be able to come up with an input program that takes a long time. + +If you'd like to file a bug, run with envvar XLA_FLAGS=--xla_dump_to=/tmp/foo and attach the results. +2025-03-24 17:03:07.590940: E external/xla/xla/service/slow_operation_alarm.cc:140] The operation took 1.033086561s +Constant folding an instruction is taking > 1s: + + %gather.423 = f32[10,107,5994]{2,1,0} gather(f32[11,107,5994]{2,1,0} %constant.7504, s32[10,1]{1,0} %constant.7532), offset_dims={1,2}, collapsed_slice_dims={0}, start_index_map={0}, index_vector_dim=1, slice_sizes={1,107,5994}, metadata={op_name="jit()/jit(main)/jit(interp2d)/jit(fun)/jit(_take)/gather" source_file="/home/annalena/rubix/rubix/core/ifu.py" source_line=107} + +This isn't necessarily a bug; constant-folding is inherently a trade-off between compilation time and speed at runtime. XLA has some guards that attempt to keep constant folding from taking too long, but fundamentally you'll always be able to come up with an input program that takes a long time. + +If you'd like to file a bug, run with envvar XLA_FLAGS=--xla_dump_to=/tmp/foo and attach the results. + Executing: 29%|██▉ | 5/17 [07:36<19:23, 96.97s/cell] Executing: 41%|████ | 7/17 [07:57<10:44, 64.40s/cell] Executing: 53%|█████▎ | 9/17 [08:35<06:27, 48.43s/cell] Executing: 59%|█████▉ | 10/17 [08:35<04:31, 38.79s/cell] Executing: 65%|██████▍ | 11/17 [08:37<03:02, 30.43s/cell] Executing: 71%|███████ | 12/17 [08:37<01:55, 23.16s/cell] Executing: 76%|███████▋ | 13/17 [08:37<01:09, 17.31s/cell] Executing: 82%|████████▏ | 14/17 [08:38<00:38, 12.88s/cell] Executing: 88%|████████▊ | 15/17 [08:39<00:19, 9.55s/cell] Executing: 100%|██████████| 17/17 [08:40<00:00, 5.51s/cell] Executing: 100%|██████████| 17/17 [08:48<00:00, 31.10s/cell] +Processing galaxy_id: g8.13e11 +Running Papermill for pipeline: calc_ifu with galaxy_id: g8.13e11 at telescope: MUSE_ultraWFM +Input Notebook: fits_file_nihao_script.ipynb +Output Notebook: fits_file_nihao_script_executed.ipynb +Passed unknown parameter: galaxy_id +Passed unknown parameter: pipeline_name +Passed unknown parameter: telescope_name +Input notebook does not contain a cell with tag 'parameters' + Executing: 0%| | 0/17 [00:00 1s: + + %gather.1049 = f32[10,107,5994]{2,1,0} gather(f32[11,107,5994]{2,1,0} %constant.15554, s32[10,1]{1,0} %constant.15583), offset_dims={1,2}, collapsed_slice_dims={0}, start_index_map={0}, index_vector_dim=1, slice_sizes={1,107,5994}, metadata={op_name="jit()/jit(main)/jit(interp2d)/jit(fun)/jit(_take)/gather" source_file="/home/annalena/rubix/rubix/core/ifu.py" source_line=107} + +This isn't necessarily a bug; constant-folding is inherently a trade-off between compilation time and speed at runtime. XLA has some guards that attempt to keep constant folding from taking too long, but fundamentally you'll always be able to come up with an input program that takes a long time. + +If you'd like to file a bug, run with envvar XLA_FLAGS=--xla_dump_to=/tmp/foo and attach the results. +2025-03-24 17:44:37.163503: E external/xla/xla/service/slow_operation_alarm.cc:140] The operation took 1.003593605s +Constant folding an instruction is taking > 1s: + + %gather.1049 = f32[10,107,5994]{2,1,0} gather(f32[11,107,5994]{2,1,0} %constant.15554, s32[10,1]{1,0} %constant.15583), offset_dims={1,2}, collapsed_slice_dims={0}, start_index_map={0}, index_vector_dim=1, slice_sizes={1,107,5994}, metadata={op_name="jit()/jit(main)/jit(interp2d)/jit(fun)/jit(_take)/gather" source_file="/home/annalena/rubix/rubix/core/ifu.py" source_line=107} + +This isn't necessarily a bug; constant-folding is inherently a trade-off between compilation time and speed at runtime. XLA has some guards that attempt to keep constant folding from taking too long, but fundamentally you'll always be able to come up with an input program that takes a long time. + +If you'd like to file a bug, run with envvar XLA_FLAGS=--xla_dump_to=/tmp/foo and attach the results. + Executing: 29%|██▉ | 5/17 [18:37<47:30, 237.56s/cell] Executing: 41%|████ | 7/17 [19:25<26:08, 156.84s/cell] Executing: 53%|█████▎ | 9/17 [20:07<14:34, 109.30s/cell] Executing: 65%|██████▍ | 11/17 [20:09<07:18, 73.12s/cell] Executing: 71%|███████ | 12/17 [20:09<04:55, 59.10s/cell] Executing: 76%|███████▋ | 13/17 [20:10<03:05, 46.41s/cell] Executing: 82%|████████▏ | 14/17 [20:10<01:46, 35.60s/cell] Executing: 88%|████████▊ | 15/17 [20:11<00:53, 26.80s/cell] Executing: 94%|█████████▍| 16/17 [20:11<00:19, 19.75s/cell] Executing: 100%|██████████| 17/17 [20:12<00:00, 14.62s/cell]Wait for final termination of kernel timed out - continuing... + Executing: 100%|██████████| 17/17 [20:22<00:00, 71.94s/cell] +Processing galaxy_id: g1.92e12 +Running Papermill for pipeline: calc_ifu with galaxy_id: g1.92e12 at telescope: MUSE_ultraWFM +Input Notebook: fits_file_nihao_script.ipynb +Output Notebook: fits_file_nihao_script_executed.ipynb +Passed unknown parameter: galaxy_id +Passed unknown parameter: pipeline_name +Passed unknown parameter: telescope_name +Input notebook does not contain a cell with tag 'parameters' + Executing: 0%| | 0/17 [00:00 1s: + + %gather.687 = f32[12,105,5994]{2,1,0} gather(f32[12,106,5994]{2,1,0} %constant.13891, s32[105,1]{1,0} %constant.13923), offset_dims={0,2}, collapsed_slice_dims={1}, start_index_map={1}, index_vector_dim=1, slice_sizes={12,1,5994}, metadata={op_name="jit()/jit(main)/jit(interp2d)/jit(fun)/jit(_take)/gather" source_file="/home/annalena/rubix/rubix/core/ifu.py" source_line=107} + +This isn't necessarily a bug; constant-folding is inherently a trade-off between compilation time and speed at runtime. XLA has some guards that attempt to keep constant folding from taking too long, but fundamentally you'll always be able to come up with an input program that takes a long time. + +If you'd like to file a bug, run with envvar XLA_FLAGS=--xla_dump_to=/tmp/foo and attach the results. +2025-03-24 18:04:36.441490: E external/xla/xla/service/slow_operation_alarm.cc:140] The operation took 1.020945471s +Constant folding an instruction is taking > 1s: + + %gather.687 = f32[12,105,5994]{2,1,0} gather(f32[12,106,5994]{2,1,0} %constant.13891, s32[105,1]{1,0} %constant.13923), offset_dims={0,2}, collapsed_slice_dims={1}, start_index_map={1}, index_vector_dim=1, slice_sizes={12,1,5994}, metadata={op_name="jit()/jit(main)/jit(interp2d)/jit(fun)/jit(_take)/gather" source_file="/home/annalena/rubix/rubix/core/ifu.py" source_line=107} + +This isn't necessarily a bug; constant-folding is inherently a trade-off between compilation time and speed at runtime. XLA has some guards that attempt to keep constant folding from taking too long, but fundamentally you'll always be able to come up with an input program that takes a long time. + +If you'd like to file a bug, run with envvar XLA_FLAGS=--xla_dump_to=/tmp/foo and attach the results. + Executing: 29%|██▉ | 5/17 [23:07<59:01, 295.09s/cell] Executing: 41%|████ | 7/17 [24:08<32:30, 195.04s/cell] Executing: 53%|█████▎ | 9/17 [24:41<17:39, 132.39s/cell] Executing: 65%|██████▍ | 11/17 [24:42<08:50, 88.49s/cell] Executing: 71%|███████ | 12/17 [24:43<05:57, 71.50s/cell] Executing: 76%|███████▋ | 13/17 [24:43<03:44, 56.15s/cell] Executing: 82%|████████▏ | 14/17 [24:43<02:09, 43.06s/cell] Executing: 88%|████████▊ | 15/17 [24:44<01:04, 32.38s/cell] Executing: 100%|██████████| 17/17 [24:45<00:00, 19.02s/cell]Wait for final termination of kernel timed out - continuing... + Executing: 100%|██████████| 17/17 [24:56<00:00, 88.02s/cell] +Processing galaxy_id: g2.79e12 +Running Papermill for pipeline: calc_ifu with galaxy_id: g2.79e12 at telescope: MUSE_ultraWFM +Input Notebook: fits_file_nihao_script.ipynb +Output Notebook: fits_file_nihao_script_executed.ipynb +Passed unknown parameter: galaxy_id +Passed unknown parameter: pipeline_name +Passed unknown parameter: telescope_name +Input notebook does not contain a cell with tag 'parameters' + Executing: 0%| | 0/17 [00:00 1s: + + %gather.1241 = f32[10,107,5994]{2,1,0} gather(f32[11,107,5994]{2,1,0} %constant.19959, s32[10,1]{1,0} %constant.19989), offset_dims={1,2}, collapsed_slice_dims={0}, start_index_map={0}, index_vector_dim=1, slice_sizes={1,107,5994}, metadata={op_name="jit()/jit(main)/jit(interp2d)/jit(fun)/jit(_take)/gather" source_file="/home/annalena/rubix/rubix/core/ifu.py" source_line=107} + +This isn't necessarily a bug; constant-folding is inherently a trade-off between compilation time and speed at runtime. XLA has some guards that attempt to keep constant folding from taking too long, but fundamentally you'll always be able to come up with an input program that takes a long time. + +If you'd like to file a bug, run with envvar XLA_FLAGS=--xla_dump_to=/tmp/foo and attach the results. +2025-03-24 18:30:00.399293: E external/xla/xla/service/slow_operation_alarm.cc:140] The operation took 1.005773456s +Constant folding an instruction is taking > 1s: + + %gather.1241 = f32[10,107,5994]{2,1,0} gather(f32[11,107,5994]{2,1,0} %constant.19959, s32[10,1]{1,0} %constant.19989), offset_dims={1,2}, collapsed_slice_dims={0}, start_index_map={0}, index_vector_dim=1, slice_sizes={1,107,5994}, metadata={op_name="jit()/jit(main)/jit(interp2d)/jit(fun)/jit(_take)/gather" source_file="/home/annalena/rubix/rubix/core/ifu.py" source_line=107} + +This isn't necessarily a bug; constant-folding is inherently a trade-off between compilation time and speed at runtime. XLA has some guards that attempt to keep constant folding from taking too long, but fundamentally you'll always be able to come up with an input program that takes a long time. + +If you'd like to file a bug, run with envvar XLA_FLAGS=--xla_dump_to=/tmp/foo and attach the results. + Executing: 29%|██▉ | 5/17 [20:18<51:50, 259.18s/cell] Executing: 41%|████ | 7/17 [20:55<28:00, 168.08s/cell] Executing: 53%|█████▎ | 9/17 [21:13<14:58, 112.30s/cell] Executing: 65%|██████▍ | 11/17 [21:14<07:29, 74.97s/cell] Executing: 71%|███████ | 12/17 [21:14<05:02, 60.58s/cell] Executing: 76%|███████▋ | 13/17 [21:14<03:10, 47.55s/cell] Executing: 82%|████████▏ | 14/17 [21:14<01:49, 36.46s/cell] Executing: 88%|████████▊ | 15/17 [21:15<00:54, 27.39s/cell] Executing: 100%|██████████| 17/17 [21:16<00:00, 16.07s/cell] Executing: 100%|██████████| 17/17 [21:26<00:00, 75.65s/cell] diff --git a/notebooks/rubix_pipeline_nihao.ipynb b/notebooks/rubix_pipeline_nihao.ipynb index 9455ab75..5779757f 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": {}, @@ -46,8 +59,8 @@ " \"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", @@ -170,6 +183,105 @@ "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": [ + "# 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", + "stars_age = np.asarray(rubixdata.stars.mass)\n", + "\n", + "# Compute the sum of velocities and count per pixel using np.bincount\n", + "sum_stars_age = np.bincount(pixel_assignment, weights=stars_age[0,:])\n", + "counts = np.bincount(pixel_assignment)\n", + "\n", + "# If you know the pixel grid dimensions (for example, a square grid)\n", + "n_pixels = len(sum_stars_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", + "stars_age_map = sum_stars_age.reshape((grid_size, grid_size))\n", + "\n", + "# Plot the result\n", + "plt.figure(figsize=(6, 5))\n", + "plt.imshow(stars_age_map, origin='lower')\n", + "plt.colorbar(label='Stars Mass')\n", + "plt.title('Stars age per Pixel')\n", + "plt.xlabel('X pixel index')\n", + "plt.ylabel('Y pixel index')\n", + "#storepath = f\"/mnt/storage/annalena_data/mockNIHAO_GECKOS/datacube_NIHAO{config['data']['load_galaxy_args']['id']}_{config[\"telescope\"][\"name\"]}_{config_illustris['pipeline']['name']}_starsmass.png\"\n", + "#plt.savefig(storepath)\n", + "plt.show()\n", + "\n" + ] + }, { "cell_type": "markdown", "metadata": {}, @@ -182,7 +294,7 @@ ], "metadata": { "kernelspec": { - "display_name": "rubix", + "display_name": "Python 3", "language": "python", "name": "python3" }, @@ -196,7 +308,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/galaxy/input_handler/pynbody.py b/rubix/galaxy/input_handler/pynbody.py index fe2beec5..66c7da5f 100644 --- a/rubix/galaxy/input_handler/pynbody.py +++ b/rubix/galaxy/input_handler/pynbody.py @@ -1,12 +1,16 @@ -from .base import BaseHandler -import pynbody -import numpy as np -from rubix.utils import SFTtoAge import logging +import os + import astropy.units as u +import numpy as np +import pynbody import yaml -import os + from rubix.units import Zsun +from rubix.utils import SFTtoAge + +from .base import BaseHandler + class PynbodyHandler(BaseHandler): def __init__( @@ -84,6 +88,33 @@ 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) + self.logger.info( f"Simulation snapshot and halo data loaded successfully for classes: {load_classes}." ) From 08dbdac37a8cb0d86d4bca49cd254495fdd356d7 Mon Sep 17 00:00:00 2001 From: anschaible Date: Tue, 1 Apr 2025 15:40:43 +0200 Subject: [PATCH 02/11] fix rotation for nihao --- notebooks/rubix_pipeline_nihao.ipynb | 21 +- rubix/galaxy/alignment.py | 409 +++++++++++---------------- rubix/telescope/telescopes.yaml | 16 +- 3 files changed, 189 insertions(+), 257 deletions(-) diff --git a/notebooks/rubix_pipeline_nihao.ipynb b/notebooks/rubix_pipeline_nihao.ipynb index 5779757f..3085781f 100644 --- a/notebooks/rubix_pipeline_nihao.ipynb +++ b/notebooks/rubix_pipeline_nihao.ipynb @@ -54,7 +54,7 @@ " \"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", @@ -69,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", "}" ] @@ -271,7 +278,7 @@ "\n", "# Plot the result\n", "plt.figure(figsize=(6, 5))\n", - "plt.imshow(stars_age_map, origin='lower')\n", + "plt.imshow(stars_age_map, origin='lower', cmap='inferno')\n", "plt.colorbar(label='Stars Mass')\n", "plt.title('Stars age per Pixel')\n", "plt.xlabel('X pixel index')\n", diff --git a/rubix/galaxy/alignment.py b/rubix/galaxy/alignment.py index a568b57b..a4a5f305 100644 --- a/rubix/galaxy/alignment.py +++ b/rubix/galaxy/alignment.py @@ -1,264 +1,179 @@ -import jax.numpy as jnp -from typing import Tuple, Union -from jax.scipy.spatial.transform import Rotation -from jaxtyping import Array, Float, jaxtyped from beartype import beartype as typechecker +from jaxtyping import jaxtyped +from rubix.galaxy.alignment import rotate_galaxy as rotate_galaxy_core +from rubix.logger import get_logger -@jaxtyped(typechecker=typechecker) -def center_particles(rubixdata: object, key: str) -> object: - """ - Center the stellar particles around the galaxy center. - - Args: - rubixdata (object): The RubixData object. - key (str): The key to the particle data. - stellar_coordinates (jnp.ndarray): The coordinates of the particles. - stellar_velocities (jnp.ndarray): The velocities of the particles. - galaxy_center (jnp.ndarray): The center of the galaxy. - - Returns: - The RubixData object with the centered particles, which contain of a new set - of coordinates and velocities as jnp.ndarray. - - Example - ------- - >>> from rubix.galaxy.alignment import center_particles - >>> rubixdata = center_particles(rubixdata, "stars") - """ - if key == "stars": - particle_coordinates = rubixdata.stars.coords - particle_velocities = rubixdata.stars.velocity - elif key == "gas": - particle_coordinates = rubixdata.gas.coords - particle_velocities = rubixdata.gas.velocity - galaxy_center = rubixdata.galaxy.center - - # Check if Center is within bounds - check_bounds = ( - (galaxy_center[0] >= jnp.min(particle_coordinates[:, 0])) - & (galaxy_center[0] <= jnp.max(particle_coordinates[:, 0])) - & (galaxy_center[1] >= jnp.min(particle_coordinates[:, 1])) - & (galaxy_center[1] <= jnp.max(particle_coordinates[:, 1])) - & (galaxy_center[2] >= jnp.min(particle_coordinates[:, 2])) - & (galaxy_center[2] <= jnp.max(particle_coordinates[:, 2])) - ) - - if not check_bounds: - raise ValueError("Center is not within the bounds of the galaxy") - - # Calculate Central Velocity from median velocities within 10kpc of center - mask = jnp.linalg.norm(particle_coordinates - galaxy_center, axis=1) < 10 - # TODO this should be a median - central_velocity = jnp.median(particle_velocities[mask], axis=0) - - if key == "stars": - rubixdata.stars.coords = particle_coordinates - galaxy_center - rubixdata.stars.velocity = particle_velocities - central_velocity - elif key == "gas": - rubixdata.gas.coords = particle_coordinates - galaxy_center - rubixdata.gas.velocity = particle_velocities - central_velocity - - return rubixdata +from .data import RubixData @jaxtyped(typechecker=typechecker) -def moment_of_inertia_tensor( - positions: Float[Array, "..."], - masses: Float[Array, "..."], - halfmass_radius: Union[Float[Array, "..."], float], -) -> Float[Array, "..."]: +def get_galaxy_rotation(config: dict): """ - Calculate the moment of inertia tensor for a given set of positions and masses within the half-light radius. - Assumes the galaxy is already centered. + Get the function to rotate the galaxy based on the configuration. Args: - positions (jnp.ndarray): The positions of the particles. - masses (jnp.ndarray): The masses of the particles. - half_light_radius (float): The half-light radius of the galaxy. + config (dict): Configuration dictionary. Returns: - The moment of inertia tensor as a jnp.ndarray. + The function to rotate the galaxy. Example - ------- - >>> from rubix.galaxy.alignment import moment_of_inertia_tensor - >>> I = moment_of_inertia_tensor(rubixdata.stars.coords, rubixdata.stars.mass, rubixdata.galaxy.half_light_radius) + -------- + >>> 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) """ - distances = jnp.sqrt( - jnp.sum(positions**2, axis=1) - ) # Direct calculation since positions are already centered - - within_halfmass_radius = distances <= halfmass_radius - - # Ensure within_halfmass_radius is concrete - concrete_indices = jnp.where( - within_halfmass_radius, size=within_halfmass_radius.shape[0] - )[0] - - filtered_positions = positions[concrete_indices] - filtered_masses = masses[concrete_indices] - - I = jnp.zeros((3, 3)) - for i in range(3): - for j in range(3): - if i == j: - I = I.at[i, j].set( - jnp.sum( - filtered_masses * jnp.sum(filtered_positions**2, axis=1) - - filtered_masses * filtered_positions[:, i] ** 2 - ) - ) - else: - I = I.at[i, j].set( - -jnp.sum( - filtered_masses - * filtered_positions[:, i] - * filtered_positions[:, j] - ) + # Check if rotation information is provided under galaxy config + if "rotation" not in config["galaxy"]: + raise ValueError("Rotation information not provided in galaxy config") + + logger = get_logger() + # Check if type is provided + if "type" in config["galaxy"]["rotation"]: + # Check if type is valid: face-on or edge-on + if config["galaxy"]["rotation"]["type"] not in ["face-on", "edge-on"]: + raise ValueError("Invalid type provided in rotation information") + + # if type is face on, alpha = beta = gamma = 0 + # if type is edge on, alpha = 90, beta = gamma = 0 + if config["galaxy"]["rotation"]["type"] == "face-on": + logger.debug("Roataion Type found: Face-on") + alpha = 0.0 + beta = 0.0 + gamma = 0.0 + + else: + # type is edge-on + logger.debug("Roataion Type found: edge-on") + alpha = 90.0 + beta = 0.0 + gamma = 0.0 + + else: + # If type is not provided, then alpha, beta, gamma should be set + # Check if alpha, beta, gamma are provided + for key in ["alpha", "beta", "gamma"]: + if key not in config["galaxy"]["rotation"]: + raise ValueError(f"{key} not provided in rotation information") + + # Get the rotation angles from the user config + alpha = config["galaxy"]["rotation"]["alpha"] + beta = config["galaxy"]["rotation"]["beta"] + gamma = config["galaxy"]["rotation"]["gamma"] + + @jaxtyped(typechecker=typechecker) + 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 + coords = component.coords + velocities = component.velocity + masses = component.mass + halfmass_radius = rubixdata.galaxy.halfmassrad_stars + + assert ( + coords is not None + ), f"Coordinates not found for {particle_type}. " + assert ( + velocities is not None + ), f"Velocities not found for {particle_type}. " + assert masses is not None, f"Masses not found for {particle_type}. " + + # Rotate the galaxy + coords, velocities = rotate_galaxy_core( + positions=coords, + velocities=velocities, + positions_stars=rubixdata.stars.coords, + masses_stars=rubixdata.stars.mass, + halfmass_radius=halfmass_radius, + alpha=alpha, + beta=beta, + gamma=gamma, ) - return I - - -@jaxtyped(typechecker=typechecker) -def rotation_matrix_from_inertia_tensor(I: Float[Array, "..."]) -> Float[Array, "..."]: - """ - Calculate 3x3 rotation matrix by diagonalization of the moment of inertia tensor. - - Args: - I (jnp.ndarray): The moment of inertia tensor. - - Returns: - The rotation matrix as a jnp.ndarray. - """ - - eigen_values, eigen_vectors = jnp.linalg.eigh(I) - order = jnp.argsort(eigen_values) - rotation_matrix = eigen_vectors[:, order] - return rotation_matrix - - -@jaxtyped(typechecker=typechecker) -def apply_init_rotation( - positions: Float[Array, "..."], rotation_matrix: Float[Array, "..."] -) -> Float[Array, "..."]: - """ - Apply a rotation matrix to a set of positions. - - Args: - positions (jnp.ndarray): The positions of the particles. - rotation_matrix (jnp.ndarray): The rotation matrix. - - Returns: - The rotated positions as a jnp.ndarray. - """ - - return jnp.dot(positions, rotation_matrix) - - -@jaxtyped(typechecker=typechecker) -def euler_rotation_matrix( - alpha: float, beta: float, gamma: float -) -> Float[Array, "3 3"]: - """ - Create a 3x3 rotation matrix given Euler angles (in degrees) - - Args: - alpha (float): Rotation around the x-axis in degrees - beta (float): Rotation around the y-axis in degrees - gamma (float): Rotation around the z-axis in degrees - - Returns: - The rotation matrix as a jnp.ndarray. - """ - - # alpha = alpha/180*jnp.pi - # beta = beta/180*jnp.pi - # gamma = gamma/180*jnp.pi - - # Rotation around the x-axis - # R_x = jnp.array([ - # [1, 0, 0], - # [0, jnp.cos(alpha), -jnp.sin(alpha)], - # [0, jnp.sin(alpha), jnp.cos(alpha)] - # ]) - R_x = Rotation.from_euler("x", alpha, degrees=True) - - # Rotation around the y-axis (pitch) - # R_y = jnp.array([ - # [jnp.cos(beta), 0, jnp.sin(beta)], - # [0, 1, 0], - # [-jnp.sin(beta), 0, jnp.cos(beta)] - # ]) - R_y = Rotation.from_euler("y", beta, degrees=True) - - # Rotation around the z-axis (yaw) - # R_z = jnp.array([ - # [jnp.cos(gamma), -jnp.sin(gamma), 0], - # [jnp.sin(gamma), jnp.cos(gamma), 0], - # [0, 0, 1] - # ]) - R_z = Rotation.from_euler("z", gamma, degrees=True) - - # Combine the rotations by matrix multiplication: R = R_z * R_y * R_x - R = R_z * R_y * R_x - - return R.as_matrix() - - -@jaxtyped(typechecker=typechecker) -def apply_rotation( - positions: Float[Array, "* 3"], alpha: float, beta: float, gamma: float -) -> Float[Array, "* 3"]: - """ - Apply a rotation to a set of positions given Euler angles. - - Args: - positions (jnp.ndarray): The positions of the particles. - alpha (float): Rotation around the x-axis in degrees - beta (float): Rotation around the y-axis in degrees - gamma (float): Rotation around the z-axis in degrees - - Returns: - The rotated positions as a jnp.ndarray. - """ - - R = euler_rotation_matrix(alpha, beta, gamma) - return jnp.dot(positions, R) - - -@jaxtyped(typechecker=typechecker) -def rotate_galaxy( - positions: Float[Array, "* 3"], - velocities: Float[Array, "* 3"], - masses: Float[Array, "..."], - halfmass_radius: Union[Float[Array, "..."], float], - alpha: float, - beta: float, - gamma: float, -) -> Tuple[Float[Array, "* 3"], Float[Array, "* 3"]]: - """ - Orientate the galaxy by applying a rotation matrix to the positions of the particles. - - Args: - positions (jnp.ndarray): The positions of the particles. - velocities (jnp.ndarray): The velocities of the particles. - masses (jnp.ndarray): The masses of the particles. - halfmass_radius (float): The half-mass radius of the galaxy. - alpha (float): Rotation around the x-axis in degrees - beta (float): Rotation around the y-axis in degrees - gamma (float): Rotation around the z-axis in degrees - - 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) - return pos_final, vel_final + # Update the inputs + # rubixdata.stars.coords = coords + # rubixdata.stars.velocity = velocities + setattr(component, "coords", coords) + 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/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 From 62c3e2153b7729545d143efaf13f9b5037a9a127 Mon Sep 17 00:00:00 2001 From: anschaible Date: Tue, 1 Apr 2025 15:51:02 +0200 Subject: [PATCH 03/11] fix rotation for nihao galaxies --- notebooks/rubix_pipeline_nihao.ipynb | 25 +- rubix/core/rotation.py | 72 ++++- rubix/galaxy/alignment.py | 413 ++++++++++++++++----------- 3 files changed, 337 insertions(+), 173 deletions(-) diff --git a/notebooks/rubix_pipeline_nihao.ipynb b/notebooks/rubix_pipeline_nihao.ipynb index 3085781f..418a458a 100644 --- a/notebooks/rubix_pipeline_nihao.ipynb +++ b/notebooks/rubix_pipeline_nihao.ipynb @@ -255,37 +255,46 @@ "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", - "stars_age = np.asarray(rubixdata.stars.mass)\n", + "ages = np.asarray(rubixdata.stars.age[0, :])\n", "\n", "# Compute the sum of velocities and count per pixel using np.bincount\n", - "sum_stars_age = np.bincount(pixel_assignment, weights=stars_age[0,:])\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(sum_stars_age)\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", - "stars_age_map = sum_stars_age.reshape((grid_size, grid_size))\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(stars_age_map, origin='lower', cmap='inferno')\n", - "plt.colorbar(label='Stars Mass')\n", - "plt.title('Stars age per Pixel')\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\"/mnt/storage/annalena_data/mockNIHAO_GECKOS/datacube_NIHAO{config['data']['load_galaxy_args']['id']}_{config[\"telescope\"][\"name\"]}_{config_illustris['pipeline']['name']}_starsmass.png\"\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" ] }, diff --git a/rubix/core/rotation.py b/rubix/core/rotation.py index c5d68b01..a4a5f305 100644 --- a/rubix/core/rotation.py +++ b/rubix/core/rotation.py @@ -1,8 +1,10 @@ -from rubix.logger import get_logger +from beartype import beartype as typechecker +from jaxtyping import jaxtyped + from rubix.galaxy.alignment import rotate_galaxy as rotate_galaxy_core +from rubix.logger import get_logger + from .data import RubixData -from jaxtyping import jaxtyped -from beartype import beartype as typechecker @jaxtyped(typechecker=typechecker) @@ -74,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 @@ -97,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, @@ -110,6 +115,65 @@ def rotate_galaxy(rubixdata: RubixData) -> RubixData: setattr(component, "coords", coords) 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 a4a5f305..a471cbe6 100644 --- a/rubix/galaxy/alignment.py +++ b/rubix/galaxy/alignment.py @@ -1,179 +1,270 @@ +from typing import Tuple, Union + +import jax.numpy as jnp from beartype import beartype as typechecker -from jaxtyping import jaxtyped +from jax.scipy.spatial.transform import Rotation +from jaxtyping import Array, Float, jaxtyped + + +@jaxtyped(typechecker=typechecker) +def center_particles(rubixdata: object, key: str) -> object: + """ + Center the stellar particles around the galaxy center. -from rubix.galaxy.alignment import rotate_galaxy as rotate_galaxy_core -from rubix.logger import get_logger + Args: + rubixdata (object): The RubixData object. + key (str): The key to the particle data. + stellar_coordinates (jnp.ndarray): The coordinates of the particles. + stellar_velocities (jnp.ndarray): The velocities of the particles. + galaxy_center (jnp.ndarray): The center of the galaxy. + + Returns: + The RubixData object with the centered particles, which contain of a new set + of coordinates and velocities as jnp.ndarray. -from .data import RubixData + Example + ------- + >>> from rubix.galaxy.alignment import center_particles + >>> rubixdata = center_particles(rubixdata, "stars") + """ + if key == "stars": + particle_coordinates = rubixdata.stars.coords + particle_velocities = rubixdata.stars.velocity + elif key == "gas": + particle_coordinates = rubixdata.gas.coords + particle_velocities = rubixdata.gas.velocity + galaxy_center = rubixdata.galaxy.center + + # Check if Center is within bounds + check_bounds = ( + (galaxy_center[0] >= jnp.min(particle_coordinates[:, 0])) + & (galaxy_center[0] <= jnp.max(particle_coordinates[:, 0])) + & (galaxy_center[1] >= jnp.min(particle_coordinates[:, 1])) + & (galaxy_center[1] <= jnp.max(particle_coordinates[:, 1])) + & (galaxy_center[2] >= jnp.min(particle_coordinates[:, 2])) + & (galaxy_center[2] <= jnp.max(particle_coordinates[:, 2])) + ) + + if not check_bounds: + raise ValueError("Center is not within the bounds of the galaxy") + + # Calculate Central Velocity from median velocities within 10kpc of center + mask = jnp.linalg.norm(particle_coordinates - galaxy_center, axis=1) < 10 + # TODO this should be a median + central_velocity = jnp.median(particle_velocities[mask], axis=0) + + if key == "stars": + rubixdata.stars.coords = particle_coordinates - galaxy_center + rubixdata.stars.velocity = particle_velocities - central_velocity + elif key == "gas": + rubixdata.gas.coords = particle_coordinates - galaxy_center + rubixdata.gas.velocity = particle_velocities - central_velocity + + return rubixdata @jaxtyped(typechecker=typechecker) -def get_galaxy_rotation(config: dict): +def moment_of_inertia_tensor( + positions: Float[Array, "..."], + masses: Float[Array, "..."], + halfmass_radius: Union[Float[Array, "..."], float], +) -> Float[Array, "..."]: """ - Get the function to rotate the galaxy based on the configuration. + Calculate the moment of inertia tensor for a given set of positions and masses within the half-light radius. + Assumes the galaxy is already centered. Args: - config (dict): Configuration dictionary. + positions (jnp.ndarray): The positions of the particles. + masses (jnp.ndarray): The masses of the particles. + half_light_radius (float): The half-light radius of the galaxy. Returns: - The function to rotate the galaxy. + The moment of inertia tensor as a jnp.ndarray. 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) + ------- + >>> from rubix.galaxy.alignment import moment_of_inertia_tensor + >>> I = moment_of_inertia_tensor(rubixdata.stars.coords, rubixdata.stars.mass, rubixdata.galaxy.half_light_radius) """ - # Check if rotation information is provided under galaxy config - if "rotation" not in config["galaxy"]: - raise ValueError("Rotation information not provided in galaxy config") - - logger = get_logger() - # Check if type is provided - if "type" in config["galaxy"]["rotation"]: - # Check if type is valid: face-on or edge-on - if config["galaxy"]["rotation"]["type"] not in ["face-on", "edge-on"]: - raise ValueError("Invalid type provided in rotation information") - - # if type is face on, alpha = beta = gamma = 0 - # if type is edge on, alpha = 90, beta = gamma = 0 - if config["galaxy"]["rotation"]["type"] == "face-on": - logger.debug("Roataion Type found: Face-on") - alpha = 0.0 - beta = 0.0 - gamma = 0.0 - - else: - # type is edge-on - logger.debug("Roataion Type found: edge-on") - alpha = 90.0 - beta = 0.0 - gamma = 0.0 + distances = jnp.sqrt( + jnp.sum(positions**2, axis=1) + ) # Direct calculation since positions are already centered - else: - # If type is not provided, then alpha, beta, gamma should be set - # Check if alpha, beta, gamma are provided - for key in ["alpha", "beta", "gamma"]: - if key not in config["galaxy"]["rotation"]: - raise ValueError(f"{key} not provided in rotation information") - - # Get the rotation angles from the user config - alpha = config["galaxy"]["rotation"]["alpha"] - beta = config["galaxy"]["rotation"]["beta"] - gamma = config["galaxy"]["rotation"]["gamma"] - - @jaxtyped(typechecker=typechecker) - 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 - coords = component.coords - velocities = component.velocity - masses = component.mass - halfmass_radius = rubixdata.galaxy.halfmassrad_stars - - assert ( - coords is not None - ), f"Coordinates not found for {particle_type}. " - assert ( - velocities is not None - ), f"Velocities not found for {particle_type}. " - assert masses is not None, f"Masses not found for {particle_type}. " - - # Rotate the galaxy - coords, velocities = rotate_galaxy_core( - positions=coords, - velocities=velocities, - positions_stars=rubixdata.stars.coords, - masses_stars=rubixdata.stars.mass, - halfmass_radius=halfmass_radius, - alpha=alpha, - beta=beta, - gamma=gamma, + within_halfmass_radius = distances <= halfmass_radius + + # Ensure within_halfmass_radius is concrete + concrete_indices = jnp.where( + within_halfmass_radius, size=within_halfmass_radius.shape[0] + )[0] + + filtered_positions = positions[concrete_indices] + filtered_masses = masses[concrete_indices] + + I = jnp.zeros((3, 3)) + for i in range(3): + for j in range(3): + if i == j: + I = I.at[i, j].set( + jnp.sum( + filtered_masses * jnp.sum(filtered_positions**2, axis=1) + - filtered_masses * filtered_positions[:, i] ** 2 + ) ) + else: + I = I.at[i, j].set( + -jnp.sum( + filtered_masses + * filtered_positions[:, i] + * filtered_positions[:, j] + ) + ) + return I + + +@jaxtyped(typechecker=typechecker) +def rotation_matrix_from_inertia_tensor(I: Float[Array, "..."]) -> Float[Array, "..."]: + """ + Calculate 3x3 rotation matrix by diagonalization of the moment of inertia tensor. + + Args: + I (jnp.ndarray): The moment of inertia tensor. + + Returns: + The rotation matrix as a jnp.ndarray. + """ + + eigen_values, eigen_vectors = jnp.linalg.eigh(I) + order = jnp.argsort(eigen_values) + rotation_matrix = eigen_vectors[:, order] + return rotation_matrix + + +@jaxtyped(typechecker=typechecker) +def apply_init_rotation( + positions: Float[Array, "..."], rotation_matrix: Float[Array, "..."] +) -> Float[Array, "..."]: + """ + Apply a rotation matrix to a set of positions. + + Args: + positions (jnp.ndarray): The positions of the particles. + rotation_matrix (jnp.ndarray): The rotation matrix. + + Returns: + The rotated positions as a jnp.ndarray. + """ + + return jnp.dot(positions, rotation_matrix) + + +@jaxtyped(typechecker=typechecker) +def euler_rotation_matrix( + alpha: float, beta: float, gamma: float +) -> Float[Array, "3 3"]: + """ + Create a 3x3 rotation matrix given Euler angles (in degrees) + + Args: + alpha (float): Rotation around the x-axis in degrees + beta (float): Rotation around the y-axis in degrees + gamma (float): Rotation around the z-axis in degrees + + Returns: + The rotation matrix as a jnp.ndarray. + """ + + # alpha = alpha/180*jnp.pi + # beta = beta/180*jnp.pi + # gamma = gamma/180*jnp.pi + + # Rotation around the x-axis + # R_x = jnp.array([ + # [1, 0, 0], + # [0, jnp.cos(alpha), -jnp.sin(alpha)], + # [0, jnp.sin(alpha), jnp.cos(alpha)] + # ]) + R_x = Rotation.from_euler("x", alpha, degrees=True) + + # Rotation around the y-axis (pitch) + # R_y = jnp.array([ + # [jnp.cos(beta), 0, jnp.sin(beta)], + # [0, 1, 0], + # [-jnp.sin(beta), 0, jnp.cos(beta)] + # ]) + R_y = Rotation.from_euler("y", beta, degrees=True) + + # Rotation around the z-axis (yaw) + # R_z = jnp.array([ + # [jnp.cos(gamma), -jnp.sin(gamma), 0], + # [jnp.sin(gamma), jnp.cos(gamma), 0], + # [0, 0, 1] + # ]) + R_z = Rotation.from_euler("z", gamma, degrees=True) + + # Combine the rotations by matrix multiplication: R = R_z * R_y * R_x + R = R_z * R_y * R_x + + return R.as_matrix() + + +@jaxtyped(typechecker=typechecker) +def apply_rotation( + positions: Float[Array, "* 3"], alpha: float, beta: float, gamma: float +) -> Float[Array, "* 3"]: + """ + Apply a rotation to a set of positions given Euler angles. + + Args: + positions (jnp.ndarray): The positions of the particles. + alpha (float): Rotation around the x-axis in degrees + beta (float): Rotation around the y-axis in degrees + gamma (float): Rotation around the z-axis in degrees + + Returns: + The rotated positions as a jnp.ndarray. + """ + + R = euler_rotation_matrix(alpha, beta, gamma) + return jnp.dot(positions, R) + + +@jaxtyped(typechecker=typechecker) +def rotate_galaxy( + positions: Float[Array, "* 3"], + velocities: Float[Array, "* 3"], + 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. + + Args: + positions (jnp.ndarray): The positions of the particles. + velocities (jnp.ndarray): The velocities of the particles. + masses (jnp.ndarray): The masses of the particles. + halfmass_radius (float): The half-mass radius of the galaxy. + alpha (float): Rotation around the x-axis in degrees + beta (float): Rotation around the y-axis in degrees + gamma (float): Rotation around the z-axis in degrees + + Returns: + The rotated positions and velocities as a jnp.ndarray. + """ + 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) + else: + pos_final = apply_rotation(positions, alpha, beta, gamma) + vel_final = apply_rotation(velocities, alpha, beta, gamma) - # Update the inputs - # rubixdata.stars.coords = coords - # rubixdata.stars.velocity = velocities - setattr(component, "coords", coords) - 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 + return pos_final, vel_final From 9c60c7c5a3bfa67812e0628ce04f59776c1332fc Mon Sep 17 00:00:00 2001 From: anschaible Date: Tue, 1 Apr 2025 16:23:21 +0200 Subject: [PATCH 04/11] fix wrong age for nihao --- notebooks/rubix_pipeline_nihao.ipynb | 2 +- rubix/galaxy/input_handler/pynbody.py | 2 ++ 2 files changed, 3 insertions(+), 1 deletion(-) diff --git a/notebooks/rubix_pipeline_nihao.ipynb b/notebooks/rubix_pipeline_nihao.ipynb index 418a458a..17462920 100644 --- a/notebooks/rubix_pipeline_nihao.ipynb +++ b/notebooks/rubix_pipeline_nihao.ipynb @@ -139,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", diff --git a/rubix/galaxy/input_handler/pynbody.py b/rubix/galaxy/input_handler/pynbody.py index 66c7da5f..97eb4bbc 100644 --- a/rubix/galaxy/input_handler/pynbody.py +++ b/rubix/galaxy/input_handler/pynbody.py @@ -115,6 +115,8 @@ def load_data(self): self.logger.info("Metals assigned to gas particles.") self.logger.info("Metals shape is: %s", self.data["gas"]["metals"].shape) + self.data["stars"]["age"] = 14.14019767 * u.Gyr - self.data["stars"]["age"] + self.logger.info( f"Simulation snapshot and halo data loaded successfully for classes: {load_classes}." ) From d5e303f510a528d11014a72fad30b4f6ecd8c809 Mon Sep 17 00:00:00 2001 From: anschaible Date: Tue, 1 Apr 2025 16:53:55 +0200 Subject: [PATCH 05/11] fits file saving --- rubix/core/fits.py | 25 +++++++++++++++++-------- 1 file changed, 17 insertions(+), 8 deletions(-) diff --git a/rubix/core/fits.py b/rubix/core/fits.py index 9a3dceac..a0cadb23 100644 --- a/rubix/core/fits.py +++ b/rubix/core/fits.py @@ -1,11 +1,14 @@ +import os + +import matplotlib.pyplot as plt import numpy as np from astropy.io import fits +from matplotlib.colors import LogNorm +from mpdaf.obj import Cube + from rubix.core.telescope import get_telescope from rubix.logger import get_logger -from mpdaf.obj import Cube -import matplotlib.pyplot as plt -from matplotlib.colors import LogNorm -import os + def store_fits(config, data, filepath): """ @@ -38,10 +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"]: + 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 + # For Illustris and NIHAO galaxy_id = config["data"]["load_galaxy_args"]["id"] snapshot = config["data"]["args"]["snapshot"] @@ -63,7 +71,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 @@ -89,7 +97,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) @@ -97,6 +105,7 @@ def store_fits(config, data, filepath): hdul.writeto(output_filename, overwrite=True) logger.info(f"Datacube saved to {output_filename}") + def load_fits(filepath): """ Load a FITS file and return the datacube. From 131f0ec209c57c5552437f6b2ae767f92416e0b3 Mon Sep 17 00:00:00 2001 From: anschaible Date: Wed, 2 Apr 2025 10:20:52 +0200 Subject: [PATCH 06/11] reduce chunk size for spectral assignment --- rubix/core/fits.py | 5 ++++- rubix/core/ifu.py | 2 +- 2 files changed, 5 insertions(+), 2 deletions(-) diff --git a/rubix/core/fits.py b/rubix/core/fits.py index a0cadb23..a766122d 100644 --- a/rubix/core/fits.py +++ b/rubix/core/fits.py @@ -41,7 +41,10 @@ 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"] - if 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"] diff --git a/rubix/core/ifu.py b/rubix/core/ifu.py index 38e5edb5..7ab7c68e 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 From 587cd68543462c8cfb0c7d3db2b41a45f1313e37 Mon Sep 17 00:00:00 2001 From: anschaible Date: Mon, 2 Jun 2025 18:41:25 +0200 Subject: [PATCH 07/11] fix pytest for galaxy rotation --- notebooks/nohup.out | 204 --------------------------------- rubix/core/rotation.py | 2 +- rubix/galaxy/alignment.py | 6 +- tests/test_galaxy_alignment.py | 2 +- 4 files changed, 6 insertions(+), 208 deletions(-) delete mode 100644 notebooks/nohup.out diff --git a/notebooks/nohup.out b/notebooks/nohup.out deleted file mode 100644 index b7dea1b2..00000000 --- a/notebooks/nohup.out +++ /dev/null @@ -1,204 +0,0 @@ -Processing galaxy_id: g5.02e11 -Running Papermill for pipeline: calc_ifu with galaxy_id: g5.02e11 at telescope: MUSE_ultraWFM -Input Notebook: fits_file_nihao_script.ipynb -Output Notebook: fits_file_nihao_script_executed.ipynb -Passed unknown parameter: galaxy_id -Passed unknown parameter: pipeline_name -Passed unknown parameter: telescope_name -Input notebook does not contain a cell with tag 'parameters' - Executing: 0%| | 0/17 [00:00 1s: - - %gather.423 = f32[10,107,5994]{2,1,0} gather(f32[11,107,5994]{2,1,0} %constant.7504, s32[10,1]{1,0} %constant.7532), offset_dims={1,2}, collapsed_slice_dims={0}, start_index_map={0}, index_vector_dim=1, slice_sizes={1,107,5994}, metadata={op_name="jit()/jit(main)/jit(interp2d)/jit(fun)/jit(_take)/gather" source_file="/home/annalena/rubix/rubix/core/ifu.py" source_line=107} - -This isn't necessarily a bug; constant-folding is inherently a trade-off between compilation time and speed at runtime. XLA has some guards that attempt to keep constant folding from taking too long, but fundamentally you'll always be able to come up with an input program that takes a long time. - -If you'd like to file a bug, run with envvar XLA_FLAGS=--xla_dump_to=/tmp/foo and attach the results. -2025-03-24 17:03:07.590940: E external/xla/xla/service/slow_operation_alarm.cc:140] The operation took 1.033086561s -Constant folding an instruction is taking > 1s: - - %gather.423 = f32[10,107,5994]{2,1,0} gather(f32[11,107,5994]{2,1,0} %constant.7504, s32[10,1]{1,0} %constant.7532), offset_dims={1,2}, collapsed_slice_dims={0}, start_index_map={0}, index_vector_dim=1, slice_sizes={1,107,5994}, metadata={op_name="jit()/jit(main)/jit(interp2d)/jit(fun)/jit(_take)/gather" source_file="/home/annalena/rubix/rubix/core/ifu.py" source_line=107} - -This isn't necessarily a bug; constant-folding is inherently a trade-off between compilation time and speed at runtime. XLA has some guards that attempt to keep constant folding from taking too long, but fundamentally you'll always be able to come up with an input program that takes a long time. - -If you'd like to file a bug, run with envvar XLA_FLAGS=--xla_dump_to=/tmp/foo and attach the results. - Executing: 29%|██▉ | 5/17 [07:36<19:23, 96.97s/cell] Executing: 41%|████ | 7/17 [07:57<10:44, 64.40s/cell] Executing: 53%|█████▎ | 9/17 [08:35<06:27, 48.43s/cell] Executing: 59%|█████▉ | 10/17 [08:35<04:31, 38.79s/cell] Executing: 65%|██████▍ | 11/17 [08:37<03:02, 30.43s/cell] Executing: 71%|███████ | 12/17 [08:37<01:55, 23.16s/cell] Executing: 76%|███████▋ | 13/17 [08:37<01:09, 17.31s/cell] Executing: 82%|████████▏ | 14/17 [08:38<00:38, 12.88s/cell] Executing: 88%|████████▊ | 15/17 [08:39<00:19, 9.55s/cell] Executing: 100%|██████████| 17/17 [08:40<00:00, 5.51s/cell] Executing: 100%|██████████| 17/17 [08:48<00:00, 31.10s/cell] -Processing galaxy_id: g8.13e11 -Running Papermill for pipeline: calc_ifu with galaxy_id: g8.13e11 at telescope: MUSE_ultraWFM -Input Notebook: fits_file_nihao_script.ipynb -Output Notebook: fits_file_nihao_script_executed.ipynb -Passed unknown parameter: galaxy_id -Passed unknown parameter: pipeline_name -Passed unknown parameter: telescope_name -Input notebook does not contain a cell with tag 'parameters' - Executing: 0%| | 0/17 [00:00 1s: - - %gather.1049 = f32[10,107,5994]{2,1,0} gather(f32[11,107,5994]{2,1,0} %constant.15554, s32[10,1]{1,0} %constant.15583), offset_dims={1,2}, collapsed_slice_dims={0}, start_index_map={0}, index_vector_dim=1, slice_sizes={1,107,5994}, metadata={op_name="jit()/jit(main)/jit(interp2d)/jit(fun)/jit(_take)/gather" source_file="/home/annalena/rubix/rubix/core/ifu.py" source_line=107} - -This isn't necessarily a bug; constant-folding is inherently a trade-off between compilation time and speed at runtime. XLA has some guards that attempt to keep constant folding from taking too long, but fundamentally you'll always be able to come up with an input program that takes a long time. - -If you'd like to file a bug, run with envvar XLA_FLAGS=--xla_dump_to=/tmp/foo and attach the results. -2025-03-24 17:44:37.163503: E external/xla/xla/service/slow_operation_alarm.cc:140] The operation took 1.003593605s -Constant folding an instruction is taking > 1s: - - %gather.1049 = f32[10,107,5994]{2,1,0} gather(f32[11,107,5994]{2,1,0} %constant.15554, s32[10,1]{1,0} %constant.15583), offset_dims={1,2}, collapsed_slice_dims={0}, start_index_map={0}, index_vector_dim=1, slice_sizes={1,107,5994}, metadata={op_name="jit()/jit(main)/jit(interp2d)/jit(fun)/jit(_take)/gather" source_file="/home/annalena/rubix/rubix/core/ifu.py" source_line=107} - -This isn't necessarily a bug; constant-folding is inherently a trade-off between compilation time and speed at runtime. XLA has some guards that attempt to keep constant folding from taking too long, but fundamentally you'll always be able to come up with an input program that takes a long time. - -If you'd like to file a bug, run with envvar XLA_FLAGS=--xla_dump_to=/tmp/foo and attach the results. - Executing: 29%|██▉ | 5/17 [18:37<47:30, 237.56s/cell] Executing: 41%|████ | 7/17 [19:25<26:08, 156.84s/cell] Executing: 53%|█████▎ | 9/17 [20:07<14:34, 109.30s/cell] Executing: 65%|██████▍ | 11/17 [20:09<07:18, 73.12s/cell] Executing: 71%|███████ | 12/17 [20:09<04:55, 59.10s/cell] Executing: 76%|███████▋ | 13/17 [20:10<03:05, 46.41s/cell] Executing: 82%|████████▏ | 14/17 [20:10<01:46, 35.60s/cell] Executing: 88%|████████▊ | 15/17 [20:11<00:53, 26.80s/cell] Executing: 94%|█████████▍| 16/17 [20:11<00:19, 19.75s/cell] Executing: 100%|██████████| 17/17 [20:12<00:00, 14.62s/cell]Wait for final termination of kernel timed out - continuing... - Executing: 100%|██████████| 17/17 [20:22<00:00, 71.94s/cell] -Processing galaxy_id: g1.92e12 -Running Papermill for pipeline: calc_ifu with galaxy_id: g1.92e12 at telescope: MUSE_ultraWFM -Input Notebook: fits_file_nihao_script.ipynb -Output Notebook: fits_file_nihao_script_executed.ipynb -Passed unknown parameter: galaxy_id -Passed unknown parameter: pipeline_name -Passed unknown parameter: telescope_name -Input notebook does not contain a cell with tag 'parameters' - Executing: 0%| | 0/17 [00:00 1s: - - %gather.687 = f32[12,105,5994]{2,1,0} gather(f32[12,106,5994]{2,1,0} %constant.13891, s32[105,1]{1,0} %constant.13923), offset_dims={0,2}, collapsed_slice_dims={1}, start_index_map={1}, index_vector_dim=1, slice_sizes={12,1,5994}, metadata={op_name="jit()/jit(main)/jit(interp2d)/jit(fun)/jit(_take)/gather" source_file="/home/annalena/rubix/rubix/core/ifu.py" source_line=107} - -This isn't necessarily a bug; constant-folding is inherently a trade-off between compilation time and speed at runtime. XLA has some guards that attempt to keep constant folding from taking too long, but fundamentally you'll always be able to come up with an input program that takes a long time. - -If you'd like to file a bug, run with envvar XLA_FLAGS=--xla_dump_to=/tmp/foo and attach the results. -2025-03-24 18:04:36.441490: E external/xla/xla/service/slow_operation_alarm.cc:140] The operation took 1.020945471s -Constant folding an instruction is taking > 1s: - - %gather.687 = f32[12,105,5994]{2,1,0} gather(f32[12,106,5994]{2,1,0} %constant.13891, s32[105,1]{1,0} %constant.13923), offset_dims={0,2}, collapsed_slice_dims={1}, start_index_map={1}, index_vector_dim=1, slice_sizes={12,1,5994}, metadata={op_name="jit()/jit(main)/jit(interp2d)/jit(fun)/jit(_take)/gather" source_file="/home/annalena/rubix/rubix/core/ifu.py" source_line=107} - -This isn't necessarily a bug; constant-folding is inherently a trade-off between compilation time and speed at runtime. XLA has some guards that attempt to keep constant folding from taking too long, but fundamentally you'll always be able to come up with an input program that takes a long time. - -If you'd like to file a bug, run with envvar XLA_FLAGS=--xla_dump_to=/tmp/foo and attach the results. - Executing: 29%|██▉ | 5/17 [23:07<59:01, 295.09s/cell] Executing: 41%|████ | 7/17 [24:08<32:30, 195.04s/cell] Executing: 53%|█████▎ | 9/17 [24:41<17:39, 132.39s/cell] Executing: 65%|██████▍ | 11/17 [24:42<08:50, 88.49s/cell] Executing: 71%|███████ | 12/17 [24:43<05:57, 71.50s/cell] Executing: 76%|███████▋ | 13/17 [24:43<03:44, 56.15s/cell] Executing: 82%|████████▏ | 14/17 [24:43<02:09, 43.06s/cell] Executing: 88%|████████▊ | 15/17 [24:44<01:04, 32.38s/cell] Executing: 100%|██████████| 17/17 [24:45<00:00, 19.02s/cell]Wait for final termination of kernel timed out - continuing... - Executing: 100%|██████████| 17/17 [24:56<00:00, 88.02s/cell] -Processing galaxy_id: g2.79e12 -Running Papermill for pipeline: calc_ifu with galaxy_id: g2.79e12 at telescope: MUSE_ultraWFM -Input Notebook: fits_file_nihao_script.ipynb -Output Notebook: fits_file_nihao_script_executed.ipynb -Passed unknown parameter: galaxy_id -Passed unknown parameter: pipeline_name -Passed unknown parameter: telescope_name -Input notebook does not contain a cell with tag 'parameters' - Executing: 0%| | 0/17 [00:00 1s: - - %gather.1241 = f32[10,107,5994]{2,1,0} gather(f32[11,107,5994]{2,1,0} %constant.19959, s32[10,1]{1,0} %constant.19989), offset_dims={1,2}, collapsed_slice_dims={0}, start_index_map={0}, index_vector_dim=1, slice_sizes={1,107,5994}, metadata={op_name="jit()/jit(main)/jit(interp2d)/jit(fun)/jit(_take)/gather" source_file="/home/annalena/rubix/rubix/core/ifu.py" source_line=107} - -This isn't necessarily a bug; constant-folding is inherently a trade-off between compilation time and speed at runtime. XLA has some guards that attempt to keep constant folding from taking too long, but fundamentally you'll always be able to come up with an input program that takes a long time. - -If you'd like to file a bug, run with envvar XLA_FLAGS=--xla_dump_to=/tmp/foo and attach the results. -2025-03-24 18:30:00.399293: E external/xla/xla/service/slow_operation_alarm.cc:140] The operation took 1.005773456s -Constant folding an instruction is taking > 1s: - - %gather.1241 = f32[10,107,5994]{2,1,0} gather(f32[11,107,5994]{2,1,0} %constant.19959, s32[10,1]{1,0} %constant.19989), offset_dims={1,2}, collapsed_slice_dims={0}, start_index_map={0}, index_vector_dim=1, slice_sizes={1,107,5994}, metadata={op_name="jit()/jit(main)/jit(interp2d)/jit(fun)/jit(_take)/gather" source_file="/home/annalena/rubix/rubix/core/ifu.py" source_line=107} - -This isn't necessarily a bug; constant-folding is inherently a trade-off between compilation time and speed at runtime. XLA has some guards that attempt to keep constant folding from taking too long, but fundamentally you'll always be able to come up with an input program that takes a long time. - -If you'd like to file a bug, run with envvar XLA_FLAGS=--xla_dump_to=/tmp/foo and attach the results. - Executing: 29%|██▉ | 5/17 [20:18<51:50, 259.18s/cell] Executing: 41%|████ | 7/17 [20:55<28:00, 168.08s/cell] Executing: 53%|█████▎ | 9/17 [21:13<14:58, 112.30s/cell] Executing: 65%|██████▍ | 11/17 [21:14<07:29, 74.97s/cell] Executing: 71%|███████ | 12/17 [21:14<05:02, 60.58s/cell] Executing: 76%|███████▋ | 13/17 [21:14<03:10, 47.55s/cell] Executing: 82%|████████▏ | 14/17 [21:14<01:49, 36.46s/cell] Executing: 88%|████████▊ | 15/17 [21:15<00:54, 27.39s/cell] Executing: 100%|██████████| 17/17 [21:16<00:00, 16.07s/cell] Executing: 100%|██████████| 17/17 [21:26<00:00, 75.65s/cell] diff --git a/rubix/core/rotation.py b/rubix/core/rotation.py index a4a5f305..3a0c0796 100644 --- a/rubix/core/rotation.py +++ b/rubix/core/rotation.py @@ -115,7 +115,7 @@ def rotate_galaxy(rubixdata: RubixData) -> RubixData: setattr(component, "coords", coords) setattr(component, "velocity", velocities) - return rubixdata" + return rubixdata """ logger.info("Rotating galaxy for simulation: " + config["simulation"]["name"]) # Rotate gas diff --git a/rubix/galaxy/alignment.py b/rubix/galaxy/alignment.py index a471cbe6..76e18b29 100644 --- a/rubix/galaxy/alignment.py +++ b/rubix/galaxy/alignment.py @@ -256,15 +256,17 @@ def rotate_galaxy( Returns: The rotated positions and velocities as a jnp.ndarray. """ - if key == "ILlustrisTNG": + 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) - else: + 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/tests/test_galaxy_alignment.py b/tests/test_galaxy_alignment.py index d8af9e84..ecc33036 100644 --- a/tests/test_galaxy_alignment.py +++ b/tests/test_galaxy_alignment.py @@ -186,7 +186,7 @@ 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 From 11048d88ee0438a577c867dbe4732efaf9cbf0a3 Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Mon, 2 Jun 2025 16:41:43 +0000 Subject: [PATCH 08/11] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- rubix/galaxy/alignment.py | 4 +++- tests/test_galaxy_alignment.py | 10 +++++++++- 2 files changed, 12 insertions(+), 2 deletions(-) diff --git a/rubix/galaxy/alignment.py b/rubix/galaxy/alignment.py index 76e18b29..6c3b57b7 100644 --- a/rubix/galaxy/alignment.py +++ b/rubix/galaxy/alignment.py @@ -267,6 +267,8 @@ def rotate_galaxy( 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'.") + raise ValueError( + f"Unknown key: {key} for the rotation. Supported keys are 'IllustrisTNG' and 'NIHAO'." + ) return pos_final, vel_final diff --git a/tests/test_galaxy_alignment.py b/tests/test_galaxy_alignment.py index ecc33036..ef2c4b69 100644 --- a/tests/test_galaxy_alignment.py +++ b/tests/test_galaxy_alignment.py @@ -186,7 +186,15 @@ def test_rotate_galaxy(): gamma = 0.0 rotated_positions, rotated_velocities = rotate_galaxy( - positions, velocities, positions, masses, halfmass_radius, alpha, beta, gamma, "IllustrisTNG" + positions, + velocities, + positions, + masses, + halfmass_radius, + alpha, + beta, + gamma, + "IllustrisTNG", ) assert rotated_positions.shape == positions.shape From d15ef5897f9eeb4ff8acdc9314d73faf6f371728 Mon Sep 17 00:00:00 2001 From: anschaible Date: Wed, 11 Jun 2025 11:33:16 +0200 Subject: [PATCH 09/11] update stellar age in pynbody --- .gitignore | 1 + rubix/galaxy/input_handler/pynbody.py | 5 +++-- 2 files changed, 4 insertions(+), 2 deletions(-) diff --git a/.gitignore b/.gitignore index cf808cb2..436c58f8 100644 --- a/.gitignore +++ b/.gitignore @@ -175,3 +175,4 @@ rubix/spectra/ssp/templates/fsps.h5 notebooks/frames notebooks/frames/* +notebooks/data/* \ No newline at end of file diff --git a/rubix/galaxy/input_handler/pynbody.py b/rubix/galaxy/input_handler/pynbody.py index 97eb4bbc..86ed3ba6 100644 --- a/rubix/galaxy/input_handler/pynbody.py +++ b/rubix/galaxy/input_handler/pynbody.py @@ -8,7 +8,7 @@ from rubix.units import Zsun from rubix.utils import SFTtoAge - +from rubix.cosmology import PLANCK15 as rubix_cosmo from .base import BaseHandler @@ -115,7 +115,8 @@ def load_data(self): self.logger.info("Metals assigned to gas particles.") self.logger.info("Metals shape is: %s", self.data["gas"]["metals"].shape) - self.data["stars"]["age"] = 14.14019767 * u.Gyr - self.data["stars"]["age"] + 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}." From a64b863e90496321b91177ed7525032c9d43c760 Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Wed, 11 Jun 2025 09:36:23 +0000 Subject: [PATCH 10/11] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- .gitignore | 2 +- rubix/galaxy/input_handler/pynbody.py | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/.gitignore b/.gitignore index 436c58f8..c091f833 100644 --- a/.gitignore +++ b/.gitignore @@ -175,4 +175,4 @@ rubix/spectra/ssp/templates/fsps.h5 notebooks/frames notebooks/frames/* -notebooks/data/* \ No newline at end of file +notebooks/data/* diff --git a/rubix/galaxy/input_handler/pynbody.py b/rubix/galaxy/input_handler/pynbody.py index 1a97ce31..1fc1f29e 100644 --- a/rubix/galaxy/input_handler/pynbody.py +++ b/rubix/galaxy/input_handler/pynbody.py @@ -6,9 +6,9 @@ import pynbody import yaml +from rubix.cosmology import PLANCK15 as rubix_cosmo from rubix.units import Zsun from rubix.utils import SFTtoAge -from rubix.cosmology import PLANCK15 as rubix_cosmo from .base import BaseHandler From 80389017f511474f08a9aa2863a428310284565d Mon Sep 17 00:00:00 2001 From: anschaible Date: Wed, 2 Jul 2025 16:41:04 +0200 Subject: [PATCH 11/11] add comment in rotation --- rubix/galaxy/alignment.py | 6 ++++++ 1 file changed, 6 insertions(+) diff --git a/rubix/galaxy/alignment.py b/rubix/galaxy/alignment.py index 6c3b57b7..09b41538 100644 --- a/rubix/galaxy/alignment.py +++ b/rubix/galaxy/alignment.py @@ -256,6 +256,12 @@ def rotate_galaxy( Returns: The rotated positions and velocities as a jnp.ndarray. """ + # 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)