diffsky 0.2.1__tar.gz → 0.2.2__tar.gz

{diffsky-0.2.1 → diffsky-0.2.2}/CHANGES.rst

@@ -1,3 +1,8 @@
+0.2.2 (unreleased)
+-------------------
+- Change little h convention of Monte Carlo halo generators (https://github.com/ArgonneCPAC/diffsky/pull/184)
+
+
 0.2.1 (2025-06-30)
 -------------------
 - Incorporate synthetic lightcone into mock production pipeline (https://github.com/ArgonneCPAC/diffsky/pull/165)

{diffsky-0.2.1 → diffsky-0.2.2}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: diffsky
-Version: 0.2.1
+Version: 0.2.2
 Summary: Library for differentiable generation of synthetic skies
 Author-email: Andrew Hearin <ahearin@anl.gov>
 License: BSD 3-Clause License

diffsky-0.2.2/diffsky/_version.py

@@ -0,0 +1 @@
+__version__ = '0.2.2'

{diffsky-0.2.1 → diffsky-0.2.2}/diffsky/data_loaders/hacc_utils/data_validation/tests/test_validate_lc_cores.py

@@ -87,3 +87,13 @@ def test_check_top_host_idx_tag_agreement():
         if len(msg) > 0:
             bn = os.path.basename(fn_lc_cores)
             raise ValueError(f"{bn} has mismatching top_host_idx")
+
+
+@pytest.mark.skipif(not CAN_RUN_LJ_DATA_TESTS, reason=POBOY_MSG)
+def test_check_host_pos_is_near_galaxy_pos():
+    """"""
+    bnpat = vlcc.BNPAT_LC_CORES.format("*", "*")
+    fn_list = glob(os.path.join(DRN_LC_CORES_POBOY, bnpat))
+    for fn_lc_cores in fn_list:
+        msg = vlcc.check_host_pos_is_near_galaxy_pos(fn_lc_cores)
+        assert len(msg) == 0, msg

diffsky-0.2.2/diffsky/data_loaders/hacc_utils/data_validation/tests/test_validate_lc_mock.py

@@ -0,0 +1,36 @@
+""""""
+
+import os
+from glob import glob
+
+import pytest
+
+from .. import validate_lc_mock as vlcm
+
+DRN_LC_MOCK_POBOY = "/Users/aphearin/work/DATA/LastJourney/lc_mock"
+
+try:
+    from haccytrees import Simulation as HACCSim  # noqa
+
+    HAS_HACCYTREES = True
+except ImportError:
+    HAS_HACCYTREES = False
+
+
+try:
+    assert os.path.isdir(DRN_LC_MOCK_POBOY)
+    CAN_RUN_LJ_DATA_TESTS = True
+except AssertionError:
+    CAN_RUN_LJ_DATA_TESTS = False
+CAN_RUN_LJ_DATA_TESTS = CAN_RUN_LJ_DATA_TESTS & HAS_HACCYTREES
+POBOY_MSG = "This test only runs on poboy machine with haccytrees installed"
+
+
+@pytest.mark.skipif(not CAN_RUN_LJ_DATA_TESTS, reason=POBOY_MSG)
+def test_check_all_columns_are_finite():
+    """"""
+    bnpat = vlcm.BNPAT_LC_MOCK.format("*", "*")
+    fn_list = glob(os.path.join(DRN_LC_MOCK_POBOY, bnpat))
+    for fn_lc_cores in fn_list:
+        msg = vlcm.check_all_columns_are_finite(fn_lc_cores)
+        assert len(msg) == 0

{diffsky-0.2.1 → diffsky-0.2.2}/diffsky/data_loaders/hacc_utils/data_validation/validate_lc_cores.py

@@ -102,3 +102,38 @@ def check_top_host_idx_tag_agreement(fn_lc_cores, lc_cores=None):
         msg.append(s)
 
     return msg
+
+
+def check_host_pos_is_near_galaxy_pos(fn_lc_cores, lc_cores=None):
+    """host position should be reasonably close to galaxy position"""
+    if lc_cores is None:
+        lc_cores = load_flat_hdf5(fn_lc_cores)
+
+    bn = os.path.basename(fn_lc_cores)
+
+    host_x = lc_cores["x"][lc_cores["top_host_idx"]]
+    host_y = lc_cores["y"][lc_cores["top_host_idx"]]
+    host_z = lc_cores["z"][lc_cores["top_host_idx"]]
+    dx = lc_cores["x"] - host_x
+    dy = lc_cores["y"] - host_y
+    dz = lc_cores["z"] - host_z
+    host_dist = np.sqrt(dx**2 + dy**2 + dz**2)
+
+    msg = []
+    n_very_far = np.sum(host_dist > 5)
+    if n_very_far > 10:
+        s = f"{n_very_far} galaxies in {bn} with "
+        s += "unexpectedly large xyz distance from top_host_idx"
+        msg.append(s)
+
+    msk_cen = lc_cores["central"]
+    mean_sat_dist = np.abs(np.mean(host_dist[~msk_cen]))
+    std_sat_dist = np.std(host_dist[~msk_cen])
+    if mean_sat_dist > 1:
+        s = f"<dist_sat>={mean_sat_dist:.2f} Mpc/h is unexpectedly large"
+        msg.append(s)
+    if std_sat_dist > 0.5:
+        s = f"std(dist_sat)={std_sat_dist:.2f} Mpc/h is unexpectedly large"
+        msg.append(s)
+
+    return msg

diffsky-0.2.2/diffsky/data_loaders/hacc_utils/data_validation/validate_lc_mock.py

@@ -0,0 +1,109 @@
+""" """
+
+import os
+
+import h5py
+import numpy as np
+
+from .. import load_flat_hdf5
+
+BNPAT_LC_MOCK = "data-{0}.{1}.diffsky_gals.hdf5"
+
+HLINE = "----------"
+
+
+def get_lc_mock_data_report(fn_lc_mock):
+    report = dict()
+    data = load_flat_hdf5(fn_lc_mock, dataset="data")
+
+    msg = check_all_columns_are_finite(fn_lc_mock, data=data)
+    if len(msg) > 0:
+        report["finite_colums"] = msg
+
+    msg = check_host_pos_is_near_galaxy_pos(fn_lc_mock, data=data)
+    if len(msg) > 0:
+        report["nfw_host_distance"] = msg
+
+    msg = check_all_data_columns_have_metadata(fn_lc_mock)
+    if len(msg) > 0:
+        report["column_metadata"] = msg
+
+    return report
+
+
+def write_lc_mock_report_to_disk(report, fn_lc_mock, drn_report):
+    if len(report) > 0:
+        bn_report = os.path.basename(fn_lc_mock).replace(".hdf5", ".report.txt")
+        fn_report = os.path.join(drn_report, bn_report)
+        with open(fn_report, "w") as fn_out:
+            for test_key, test_result in report.items():
+                fn_out.write(test_key + "\n")
+                for line in test_result:
+                    fn_out.write(line + "\n")
+                fn_out.write(HLINE + "\n")
+
+
+def check_all_columns_are_finite(fn_lc_mock, data=None):
+    bn = os.path.basename(fn_lc_mock)
+
+    if data is None:
+        data = load_flat_hdf5(fn_lc_mock, dataset="data")
+
+    msg = []
+    for key, arr in data.items():
+        if not np.all(np.isfinite(arr)):
+            s = f"Column {key} in {bn} has either NaN or inf"
+            msg.append(s)
+
+    return msg
+
+
+def check_all_data_columns_have_metadata(fn_lc_mock):
+
+    msg = []
+    with h5py.File(fn_lc_mock, "r") as hdf:
+        for key in hdf["data"].keys():
+            try:
+                unit = hdf["data/" + key].attrs["unit"]
+                assert len(unit) > 0
+                description = hdf["data/" + key].attrs["description"]
+                assert len(description) > 0
+            except (KeyError, AssertionError):
+                s = f"{key} is missing metadata"
+                msg.append(s)
+    return msg
+
+
+def check_host_pos_is_near_galaxy_pos(fn_lc_mock, data=None):
+    """host position should be reasonably close to galaxy position"""
+    if data is None:
+        data = load_flat_hdf5(fn_lc_mock)
+
+    bn = os.path.basename(fn_lc_mock)
+
+    host_x = data["x"][data["top_host_idx"]]
+    host_y = data["y"][data["top_host_idx"]]
+    host_z = data["z"][data["top_host_idx"]]
+    dx = data["x_nfw"] - host_x
+    dy = data["y_nfw"] - host_y
+    dz = data["z_nfw"] - host_z
+    host_dist = np.sqrt(dx**2 + dy**2 + dz**2)
+
+    msg = []
+    n_very_far = np.sum(host_dist > 5)
+    if n_very_far > 10:
+        s = f"{n_very_far} galaxies in {bn} with "
+        s += "unexpectedly large xyz distance from top_host_idx"
+        msg.append(s)
+
+    msk_cen = data["central"]
+    mean_sat_dist = np.abs(np.mean(host_dist[~msk_cen]))
+    std_sat_dist = np.std(host_dist[~msk_cen])
+    if mean_sat_dist > 1.0:
+        s = f"<dist_sat>={mean_sat_dist:.2f} Mpc/h is unexpectedly large"
+        msg.append(s)
+    if std_sat_dist > 1.0:
+        s = f"std(dist_sat)={std_sat_dist:.2f} Mpc/h is unexpectedly large"
+        msg.append(s)
+
+    return msg

{diffsky-0.2.1 → diffsky-0.2.2}/diffsky/data_loaders/hacc_utils/lc_mock_production.py

@@ -18,11 +18,43 @@ from jax import random as jran
 from ...fake_sats import halo_boundary_functions as hbf
 from ...fake_sats import nfw_config_space as nfwcs
 from ...utils.sfh_utils import get_logsm_logssfr_at_t_obs
+from . import lightcone_utils as hlu
 from . import load_lc_cf
 
 LC_CF_BNPAT = "lc_cores-{0}.{1}.diffsky_data.hdf5"
 LC_MOCK_BNPAT = LC_CF_BNPAT.replace("diffsky_data", "diffsky_gals")
 
+shapes_1 = [f"infall_fof_halo_eigS1{x}" for x in ("X", "Y", "Z")]
+shapes_2 = [f"infall_fof_halo_eigS2{x}" for x in ("X", "Y", "Z")]
+shapes_3 = [f"infall_fof_halo_eigS3{x}" for x in ("X", "Y", "Z")]
+SHAPE_KEYS = (*shapes_1, *shapes_2, *shapes_3)
+TOP_HOST_SHAPE_KEYS = ["top_host_" + key for key in SHAPE_KEYS]
+
+LC_DATA_KEYS_OUT = (
+    "core_tag",
+    "x",
+    "y",
+    "z",
+    "x_nfw",
+    "y_nfw",
+    "z_nfw",
+    "top_host_idx",
+    "central",
+    "ra_nfw",
+    "dec_nfw",
+)
+
+DIFFSKY_DATA_KEYS_OUT = (
+    "x_host",
+    "y_host",
+    "z_host",
+    "vx",
+    "vy",
+    "vz",
+    "logmp_obs_host",
+    *TOP_HOST_SHAPE_KEYS,
+)
+
 
 def write_lc_sfh_mock_to_disk(fnout, lc_data, diffsky_data):
     with h5py.File(fnout, "w") as hdf_out:
@@ -47,24 +79,24 @@ def write_lc_sfh_mock_to_disk(fnout, lc_data, diffsky_data):
         hdf_out["data/logsm_obs"] = diffsky_data["logsm_obs"]
         hdf_out["data/logssfr_obs"] = diffsky_data["logssfr_obs"]
 
-        hdf_out["data/z_true"] = lc_data["z_true"]
+        hdf_out["data/redshift_true"] = lc_data["redshift_true"]
         hdf_out["data/ra"] = lc_data["phi"]
         hdf_out["data/dec"] = np.pi / 2.0 - lc_data["theta"]
         hdf_out["data/snapnum"] = lc_data["snapnum"]
 
-        lc_data_keys_out = ("core_tag", "x", "y", "z", "top_host_idx", "central")
-        for key in lc_data_keys_out:
+        for key in LC_DATA_KEYS_OUT:
             key_out = "data/" + key
             hdf_out[key_out] = lc_data[key]
 
-        diffsky_data_keys_out = ("x_host", "y_host", "z_host", "logmp_obs_host")
-        for key in diffsky_data_keys_out:
+        for key in DIFFSKY_DATA_KEYS_OUT:
             key_out = "data/" + key
             hdf_out[key_out] = diffsky_data[key]
 
 
 def add_sfh_quantities_to_mock(sim_info, lc_data, diffsky_data, ran_key):
-    lc_data["t_obs"] = flat_wcdm.age_at_z(lc_data["z_true"], *sim_info.cosmo_params)
+    lc_data["t_obs"] = flat_wcdm.age_at_z(
+        lc_data["redshift_true"], *sim_info.cosmo_params
+    )
 
     mah_params, msk_has_diffmah_fit = load_lc_cf.get_imputed_mah_params(
         ran_key, diffsky_data, lc_data, sim_info.lgt0
@@ -132,7 +164,7 @@ def reposition_satellites(sim_info, lc_data, diffsky_data, ran_key, fixed_conc=5
     diffsky_data["y_host"] = host_pos[:, 1]
     diffsky_data["z_host"] = host_pos[:, 2]
 
-    args = (10**host_logmp_obs, sim_info.cosmo_params, lc_data["z_true"], "200m")
+    args = (10**host_logmp_obs, sim_info.cosmo_params, lc_data["redshift_true"], "200m")
     host_radius_mpc = hbf.halo_mass_to_halo_radius(*args) / 1000.0
 
     n_cores = host_logmp_obs.shape[0]
@@ -152,6 +184,10 @@ def reposition_satellites(sim_info, lc_data, diffsky_data, ran_key, fixed_conc=5
     lc_data["y_nfw"] = new_pos[:, 1]
     lc_data["z_nfw"] = new_pos[:, 2]
 
+    ra, dec = hlu.get_ra_dec(lc_data["x_nfw"], lc_data["y_nfw"], lc_data["z_nfw"])
+    lc_data["ra_nfw"] = ra
+    lc_data["dec_nfw"] = dec
+
     return lc_data, diffsky_data
 
 

{diffsky-0.2.1 → diffsky-0.2.2}/diffsky/data_loaders/hacc_utils/lightcone_utils.py

@@ -8,7 +8,9 @@ from glob import glob
 import h5py
 import numpy as np
 from diffmah.defaults import DEFAULT_MAH_PARAMS
+from dsps.cosmology import flat_wcdm
 from jax import jit as jjit
+from jax import numpy as jnp
 from jax import vmap
 
 from .. import load_flat_hdf5
@@ -23,10 +25,18 @@ SQDEG_OF_SPHERE = SQDEG_PER_STER * 4 * np.pi
 TOP_HOST_MAH_KEYS = ["top_host_" + key for key in DEFAULT_MAH_PARAMS._fields]
 SECONDARY_HOST_MAH_KEYS = ["sec_host_" + key for key in DEFAULT_MAH_PARAMS._fields]
 
+shapes_1 = [f"infall_fof_halo_eigS1{x}" for x in ("X", "Y", "Z")]
+shapes_2 = [f"infall_fof_halo_eigS2{x}" for x in ("X", "Y", "Z")]
+shapes_3 = [f"infall_fof_halo_eigS3{x}" for x in ("X", "Y", "Z")]
+SHAPE_KEYS = [*shapes_1, *shapes_2, *shapes_3]
+TOP_HOST_SHAPE_KEYS = ["top_host_" + key for key in SHAPE_KEYS]
+
+
 LC_PATCH_OUT_KEYS = (
     *DEFAULT_MAH_PARAMS._fields,
     *TOP_HOST_MAH_KEYS,
     *SECONDARY_HOST_MAH_KEYS,
+    *TOP_HOST_SHAPE_KEYS,
     "loss",
     "n_points_per_fit",
     "indx_t_ult_inf",
@@ -34,6 +44,9 @@ LC_PATCH_OUT_KEYS = (
     "n_cf_match",
     "mp_obs",
     "mp0",
+    "vx",
+    "vy",
+    "vz",
 )
 LC_PATCH_OUT_INT_KEYS = (
     "n_points_per_fit",
@@ -69,6 +82,89 @@ def jnp_take_matrix(matrix, indxarr):
     return _jnp_take_vmap(matrix, indxarr)
 
 
+@jjit
+def get_theta_phi(x, y, z):
+    """Compute lightcone angles from simulation xyz coords
+
+    Parameters
+    ----------
+    x,y,z : array, shape (n, )
+       Cartesian coords in any units
+
+    Returns
+    -------
+    theta : array, shape (n, )
+        colatitude in radians
+        0 < θ < π
+
+    phi : array, shape (n, )
+        longitude in radians
+        0 < φ < 2π
+
+    """
+    rsq = x * x + y * y + z * z
+    r = jnp.sqrt(rsq)
+    theta = jnp.arccos(z / r)
+    phi = jnp.arctan2(y, x) + jnp.pi
+    return theta, phi
+
+
+@jjit
+def get_ra_dec(x, y, z):
+    """Compute ra, dec in degrees from simulation xyz coords
+
+    Parameters
+    ----------
+    x,y,z : array, shape (n, )
+       Cartesian coords in any units
+
+    Returns
+    -------
+    ra : array, shape (n, )
+        0 < ra < 360
+
+    dec : array, shape (n, )
+        -90 < dec < 90
+
+    """
+    theta, phi = get_theta_phi(x, y, z)
+    ra, dec = _get_lon_lat_from_theta_phi(theta, phi)
+    return ra, dec
+
+
+@jjit
+def _get_lon_lat_from_theta_phi(theta, phi):
+    lon = jnp.degrees(phi)
+    lat = 90.0 - jnp.degrees(theta)
+    return lon, lat
+
+
+@jjit
+def get_redshift_from_xyz(x_mpch, y_mpch, z_mpch, cosmo_params):
+    """Compute redshift from simulation xyz coords
+
+    Parameters
+    ----------
+    x,y,z : array, shape (n, )
+       Cartesian coords in Mpc/h
+
+    Returns
+    -------
+    redshift : array, shape (n, )
+
+    """
+    d_mpch = jnp.sqrt(x_mpch**2 + y_mpch**2 + z_mpch**2)
+    d_mpc = d_mpch / cosmo_params.h
+
+    z_table = jnp.linspace(0, 20, 1_000)
+    d_table = flat_wcdm.comoving_distance(
+        z_table, cosmo_params.Om0, cosmo_params.w0, cosmo_params.wa, cosmo_params.h
+    )
+    redshift = jnp.interp(d_mpc, d_table, z_table)
+
+    return redshift
+
+
 def read_lc_ra_dec_patch_decomposition(fn):
     collector = []
     with open(fn, "r") as f:
@@ -207,6 +303,12 @@ def get_diffsky_quantities_for_lc_patch(
 
         lc_patch_data_out["n_cf_match"][msk_olap] += 1
 
+        # Matchup core velocities
+        _keys = ("vx", "vy", "vz")
+        _olap_vel = [cf_matrices[key][olap_chunk_idx][:, timestep_idx] for key in _keys]
+        for key, v in zip(_keys, _olap_vel):
+            lc_patch_data_out[key][msk_olap] = v
+
         # Matchup mah_params
         _keys = (*DEFAULT_MAH_PARAMS._fields, "loss", "n_points_per_fit")
         _olap_data = [cf_diffmah_data[key][olap_chunk_idx] for key in _keys]
@@ -215,11 +317,18 @@ def get_diffsky_quantities_for_lc_patch(
 
         # Matchup uber host_mah_params
         mah_keys = DEFAULT_MAH_PARAMS._fields
-        host_mah_keys = TOP_HOST_MAH_KEYS
-        _olap_ult_host_data = [
+        _olap_ult_host_mah_data = [
             cf_diffmah_data[key][olap_chunk_ult_host_idx] for key in mah_keys
         ]
-        for host_key, x in zip(host_mah_keys, _olap_ult_host_data):
+        for host_key, x in zip(TOP_HOST_MAH_KEYS, _olap_ult_host_mah_data):
+            lc_patch_data_out[host_key][msk_olap] = x
+
+        # Matchup uber host shapes
+        _olap_ult_host_shape_data = [
+            cf_matrices[key][:, timestep_idx][olap_chunk_ult_host_idx]
+            for key in SHAPE_KEYS
+        ]
+        for host_key, x in zip(TOP_HOST_SHAPE_KEYS, _olap_ult_host_shape_data):
             lc_patch_data_out[host_key][msk_olap] = x
 
         # Matchup secondary host_mah_params
@@ -388,3 +497,4 @@ def get_a_range_of_lc_cores_file(bname_lc_cores, sim_name):
     a_min_expected, a_max_expected = aarr[indx_step], aarr[indx_step + 1]
 
     return a_min_expected, a_max_expected
+    return a_min_expected, a_max_expected

{diffsky-0.2.1 → diffsky-0.2.2}/diffsky/data_loaders/hacc_utils/load_lc_cf.py

@@ -21,7 +21,7 @@ try:
 except ImportError:
     HAS_HACCYTREES = False
 
-SIM_INFO_KEYS = ("sim", "cosmo_params", "z_sim", "t_sim", "lgt0", "fb")
+SIM_INFO_KEYS = ("sim", "cosmo_params", "z_sim", "t_sim", "lgt0", "fb", "num_subvols")
 DiffskySimInfo = namedtuple("DiffskySimInfo", SIM_INFO_KEYS)
 
 
@@ -38,7 +38,16 @@ def get_diffsky_info_from_hacc_sim(sim_name):
     lgt0 = np.log10(t0)
     fb = sim.cosmo.Omega_b / sim.cosmo.Omega_m
 
-    diffsky_info = DiffskySimInfo(sim, cosmo_params, z_sim, t_sim, lgt0, fb)
+    if sim_name == "LastJourney":
+        num_subvols = 256
+    elif sim_name in ("DiscoveryLCDM", "DiscoveryW0WA"):
+        num_subvols = 96
+    else:
+        num_subvols = np.nan
+
+    diffsky_info = DiffskySimInfo(
+        sim, cosmo_params, z_sim, t_sim, lgt0, fb, num_subvols
+    )
 
     return diffsky_info
 
@@ -51,9 +60,9 @@ def load_lc_diffsky_patch_data(fn_lc_diffsky, indir_lc_data):
     fn_lc = os.path.join(indir_lc_data, bn_lc)
     lc_data = load_flat_hdf5(fn_lc)
 
-    lc_data["z_true"] = 1 / lc_data["scale_factor"] - 1
+    lc_data["redshift_true"] = 1 / lc_data["scale_factor"] - 1
 
-    assert lc_data["z_true"].shape[0] == diffsky_data["logm0"].shape[0]
+    assert lc_data["redshift_true"].shape[0] == diffsky_data["logm0"].shape[0]
 
     return lc_data, diffsky_data
 
@@ -82,9 +91,9 @@ def collect_lc_diffsky_data(fn_list, drn_lc_data=None):
     for key in lc_data_collector[0].keys():
         lc_data[key] = np.concatenate([x[key] for x in lc_data_collector])
 
-    lc_data["z_true"] = 1 / lc_data["scale_factor"] - 1
+    lc_data["redshift_true"] = 1 / lc_data["scale_factor"] - 1
 
-    assert lc_data["z_true"].shape[0] == diffsky_data["logm0"].shape[0]
+    assert lc_data["redshift_true"].shape[0] == diffsky_data["logm0"].shape[0]
 
     return lc_data, diffsky_data
 

{diffsky-0.2.1 → diffsky-0.2.2}/diffsky/data_loaders/hacc_utils/load_lc_cf_synthetic.py

@@ -38,10 +38,10 @@ def load_lc_diffsky_patch_data(fn_lc_cores, sim_name, ran_key, lgmp_min, lgmp_ma
     diffsky_data = mclh.mc_lightcone_host_halo_diffmah(
         *args, logmp_cutoff=11.0, lgmp_max=lgmp_max
     )
-    diffsky_data["z_true"] = diffsky_data["z_obs"]
+    diffsky_data["redshift_true"] = diffsky_data["z_obs"]
     del diffsky_data["z_obs"]
 
-    diffsky_data["top_host_idx"] = np.arange(len(diffsky_data["z_true"])).astype(int)
+    diffsky_data["top_host_idx"] = np.arange(len(diffsky_data["redshift_true"])).astype(int)
     for key in diffsky_data["mah_params"]._fields:
         diffsky_data[key] = getattr(diffsky_data["mah_params"], key)
 
@@ -55,10 +55,10 @@ def load_lc_diffsky_patch_data(fn_lc_cores, sim_name, ran_key, lgmp_min, lgmp_ma
     diffsky_data.pop("mah_params")
 
     for key in ("x", "y", "z", "x_host", "y_host", "z_host", "ra", "dec"):
-        diffsky_data[key] = np.zeros(len(diffsky_data["z_true"])) - 1.0
+        diffsky_data[key] = np.zeros(len(diffsky_data["redshift_true"])) - 1.0
 
-    ZZ = np.zeros(len(diffsky_data["z_true"]))
-    diffsky_data["core_tag"] = -np.ones(len(diffsky_data["z_true"])).astype(int)
+    ZZ = np.zeros(len(diffsky_data["redshift_true"]))
+    diffsky_data["core_tag"] = -np.ones(len(diffsky_data["redshift_true"])).astype(int)
     diffsky_data["has_diffmah_fit"] = ZZ.astype(int) + 1
     diffsky_data["n_points_per_fit"] = ZZ.astype(int) + 10_000
     diffsky_data["loss"] = ZZ + 1e-5

{diffsky-0.2.1 → diffsky-0.2.2}/diffsky/data_loaders/hacc_utils/metadata_sfh_mock.py

@@ -9,8 +9,8 @@ from . import load_lc_cf
 column_metadata = dict()
 column_metadata["central"] = ("None", "0 for satellite, 1 for central")
 column_metadata["core_tag"] = ("None", "ID of the simulated core")
-column_metadata["ra"] = ("None", "right ascension")
-column_metadata["dec"] = ("None", "declination")
+column_metadata["ra"] = ("degrees", "right ascension")
+column_metadata["dec"] = ("degrees", "declination")
 
 column_metadata["early_index"] = ("None", "Diffmah parameter for halo mass assembly")
 column_metadata["late_index"] = ("None", "Diffmah parameter for halo mass assembly")
@@ -42,19 +42,77 @@ column_metadata["logssfr_obs"] = ("1/yr ", "Log10(SFR/Mstar) at time of observat
 
 column_metadata["top_host_idx"] = ("None", "Index of host core")
 
-column_metadata["x"] = ("Mpc", "Cartesian coordinate of lightcone position")
-column_metadata["y"] = ("Mpc", "Cartesian coordinate of lightcone position")
-column_metadata["z"] = ("Mpc", "Cartesian coordinate of lightcone position")
+column_metadata["x"] = ("Mpc/h", "Comoving coordinate of lightcone position")
+column_metadata["y"] = ("Mpc/h", "Comoving coordinate of lightcone position")
+column_metadata["z"] = ("Mpc/h", "Comoving coordinate of lightcone position")
 
-column_metadata["x_host"] = ("Mpc", "Cartesian coord of host halo lightcone position")
-column_metadata["y_host"] = ("Mpc", "Cartesian coord of host halo lightcone position")
-column_metadata["z_host"] = ("Mpc", "Cartesian coord of host halo lightcone position")
+column_metadata["x_host"] = ("Mpc/h", "Comoving coord of host halo lightcone position")
+column_metadata["y_host"] = ("Mpc/h", "Comoving coord of host halo lightcone position")
+column_metadata["z_host"] = ("Mpc/h", "Comoving coord of host halo lightcone position")
 
-column_metadata["z_true"] = ("None", "True redshift")
+column_metadata["x_nfw"] = ("Mpc/h", "Comoving coord with NFW-repositioned satellites")
+column_metadata["y_nfw"] = ("Mpc/h", "Comoving coord with NFW-repositioned satellites")
+column_metadata["z_nfw"] = ("Mpc/h", "Comoving coord with NFW-repositioned satellites")
+
+column_metadata["vx"] = ("km/s", "Comoving peculiar velocity in x-direction")
+column_metadata["vy"] = ("km/s", "Comoving peculiar velocity in y-direction")
+column_metadata["vz"] = ("km/s", "Comoving peculiar velocity in z-direction")
+
+column_metadata["top_host_infall_fof_halo_eigS1X"] = (
+    "Mpc/h",
+    "x-component of first eigendirection of halo shape (unreduced intertia tensor)",
+)
+column_metadata["top_host_infall_fof_halo_eigS1Y"] = (
+    "Mpc/h",
+    "y-component of first eigendirection of halo shape (unreduced intertia tensor)",
+)
+column_metadata["top_host_infall_fof_halo_eigS1Z"] = (
+    "Mpc/h",
+    "z-component of first eigendirection of halo shape (unreduced intertia tensor)",
+)
+
+column_metadata["top_host_infall_fof_halo_eigS2X"] = (
+    "Mpc/h",
+    "x-component of second eigendirection of halo shape (unreduced intertia tensor)",
+)
+column_metadata["top_host_infall_fof_halo_eigS2Y"] = (
+    "Mpc/h",
+    "y-component of second eigendirection of halo shape (unreduced intertia tensor)",
+)
+column_metadata["top_host_infall_fof_halo_eigS2Z"] = (
+    "Mpc/h",
+    "z-component of second eigendirection of halo shape (unreduced intertia tensor)",
+)
+
+column_metadata["top_host_infall_fof_halo_eigS3X"] = (
+    "Mpc/h",
+    "x-component of third eigendirection of halo shape (unreduced intertia tensor)",
+)
+column_metadata["top_host_infall_fof_halo_eigS3Y"] = (
+    "Mpc/h",
+    "y-component of third eigendirection of halo shape (unreduced intertia tensor)",
+)
+column_metadata["top_host_infall_fof_halo_eigS3Z"] = (
+    "Mpc/h",
+    "z-component of third eigendirection of halo shape (unreduced intertia tensor)",
+)
+
+
+column_metadata["redshift_true"] = ("None", "True redshift")
+column_metadata["ra_nfw"] = (
+    "degrees",
+    "Longitude in degrees. Based on NFW-repositioned satellites.",
+)
+column_metadata["dec_nfw"] = (
+    "degrees",
+    "Latitude in degrees. Based on NFW-repositioned satellites.",
+)
 
 
 HEADER_COMMENT = """
 This file contains diffsky galaxy data.
+Each file stores mock galaxies in a thin redshift shell of a small patch of sky.
+Each column of data includes metadata with units and comments about the column.
 Contact: ahearin@anl.gov for questions.
 """
 
@@ -98,9 +156,9 @@ def append_metadata(fnout, sim_name):
             key_out = "data/" + key
             assert key_out in hdf_out.keys(), f"{key_out} is missing from {fnout}"
 
-            units, comment = val
-            hdf_out[key_out].attrs["units"] = units
-            hdf_out[key_out].attrs["comment"] = comment
+            unit, description = val
+            hdf_out[key_out].attrs["unit"] = unit
+            hdf_out[key_out].attrs["description"] = description
 
 
 def get_dependency_versions():