Fix stellar age

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,

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"

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)

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

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