diffsky 0.3.2__tar.gz → 0.3.3__tar.gz

{diffsky-0.3.2 → diffsky-0.3.3}/.github/workflows/mock_generation_test.yaml

@@ -23,7 +23,7 @@ jobs:
           remove_dotnet: true
           remove_haskell: true
           rm_cmd: "rmz"
-    
+
       - uses: actions/checkout@v6
         with:
           fetch-depth: 0
@@ -86,14 +86,14 @@ jobs:
           export DSPS_DRN="${GITHUB_WORKSPACE}/dsps_drn"
           echo "DSPS_DRN=${DSPS_DRN}" >> $GITHUB_ENV
 
-      - name: run mock generation script
+      - name: create mock with sfh_model option
         shell: bash -l {0}
         run: |
           # Create output directory
           mkdir -p ci_test_output
 
           # Run mock production script
-          mpiexec -n 1 python scripts/LJ_LC_MOCKS/make_ou26_mock.py \
+          mpiexec -n 1 python scripts/LJ_LC_MOCKS/make_ou26_mock_batch.py \
             poboy \
             0.01 \
             0.2 \
@@ -101,7 +101,6 @@ jobs:
             1 \
             ci_test_output \
             ci_test_mock \
-            -fn_u_params sfh_model \
             -sfh_model smdpl_dr1 \
             -synthetic_cores 1 \
             -lgmp_min 12.0 \
@@ -111,10 +110,51 @@ jobs:
       - name: verify script completed
         shell: bash -l {0}
         run: |
-          echo "Mock generation script completed successfully"
+          echo "sfh_model mock generated successfully"
 
       - name: validate mock
         shell: bash -l {0}
         run: |
           python scripts/LJ_LC_MOCKS/inspect_lc_mock.py ci_test_output/synthetic_cores/smdpl_dr1
 
+      - name: remove sfh_model mock
+        shell: bash -l {0}
+        run: |
+          rm -rf ci_test_output
+
+      - name: create mock with cosmos_fit option
+        shell: bash -l {0}
+        run: |
+          # Create output directory
+          mkdir -p ci_test_output
+
+          # Run mock production script
+          mpiexec -n 1 python scripts/LJ_LC_MOCKS/make_ou26_mock_batch.py \
+            poboy \
+            0.01 \
+            0.2 \
+            0 \
+            1 \
+            ci_test_output \
+            ci_test_mock \
+            -cosmos_fit cosmos260105 \
+            -synthetic_cores 1 \
+            -lgmp_min 12.0 \
+            -lgmp_max 13.0 \
+            --lsst_only
+
+      - name: verify script completed
+        shell: bash -l {0}
+        run: |
+          echo "sfh_model mock generated successfully"
+
+      - name: validate mock
+        shell: bash -l {0}
+        run: |
+          python scripts/LJ_LC_MOCKS/inspect_lc_mock.py ci_test_output/synthetic_cores/cosmos260105
+
+      - name: remove cosmos_fit mock
+        shell: bash -l {0}
+        run: |
+          rm -rf ci_test_output
+

diffsky-0.3.3/.github/workflows/shape_optimization.yaml

@@ -0,0 +1,85 @@
+name: Test script for optimizing galaxy shapes
+
+on:
+  workflow_dispatch: null
+  schedule:
+    # Runs "every Monday & Thursday at 3:05am Central"
+    - cron: '5 8 * * 1,4'
+  push:
+    branches:
+      - main
+  pull_request: null
+
+jobs:
+  mock-generation:
+    name: python fit_bulge_shapes_to_rp13.py
+    runs-on: "ubuntu-latest"
+
+    steps:
+      - name: free disk space
+        uses: endersonmenezes/free-disk-space@v3
+        with:
+          remove_android: true
+          remove_dotnet: true
+          remove_haskell: true
+          rm_cmd: "rmz"
+
+      - uses: actions/checkout@v6
+        with:
+          fetch-depth: 0
+
+      - uses: conda-incubator/setup-miniconda@v3
+        with:
+          python-version: 3.11
+          channels: conda-forge
+          channel-priority: strict
+          show-channel-urls: true
+          miniforge-version: latest
+
+      - name: configure conda and install code
+        shell: bash -l {0}
+        env:
+          GITHUB_TOKEN: ${{ secrets.CPAC_PRIVATE_REPO_TOKEN }}
+          GITHUB_USER: ${{ secrets.GH_USR_TESTING_TOKEN }}
+
+        run: |
+          conda config --set always_yes yes
+          conda install --quiet \
+            --file=requirements.txt
+          conda install --quiet \
+            pip \
+            setuptools \
+            "setuptools_scm>=7,<8" \
+            matplotlib \
+            astropy \
+            h5py \
+            mpi4py \
+            mpich
+          conda uninstall diffmah --force
+          conda uninstall diffstar --force
+          conda uninstall dsps --force
+          pip install --no-deps --no-build-isolation git+https://github.com/ArgonneCPAC/diffmah.git
+          pip install --no-deps --no-build-isolation git+https://github.com/ArgonneCPAC/diffstar.git
+          pip install --no-deps --no-build-isolation git+https://github.com/ArgonneCPAC/dsps.git
+          python -m pip install --no-build-isolation --no-deps -e .
+
+      - name: run disk optimization script
+        shell: bash -l {0}
+        run: |
+          python scripts/ellipsoid_shape_fitting/fit_disk_shapes_to_rp13.py -nsteps 20
+
+      - name: verify disk script completed
+        shell: bash -l {0}
+        run: |
+          echo "Disk optimization script completed successfully"
+
+      - name: run bulge optimization script
+        shell: bash -l {0}
+        run: |
+          python scripts/ellipsoid_shape_fitting/fit_bulge_shapes_to_rp13.py -nsteps 20
+
+      - name: verify bulge script completed
+        shell: bash -l {0}
+        run: |
+          echo "Bulge optimization script completed successfully"
+

{diffsky-0.3.2 → diffsky-0.3.3}/.github/workflows/test_main_branch.yaml

@@ -45,6 +45,7 @@ jobs:
             setuptools \
             "setuptools_scm>=7,<8" \
             matplotlib \
+            scipy \
             astropy \
             python-build
           pip uninstall diffmah --yes

{diffsky-0.3.2 → diffsky-0.3.3}/CHANGES.rst

@@ -1,3 +1,24 @@
+0.3.3 (unreleased)
+-------------------
+- Fix bug in elliptical projections
+    - https://github.com/ArgonneCPAC/diffsky/pull/297
+- Reoptimize disk/bulge shapes
+    - https://github.com/ArgonneCPAC/diffsky/pull/301
+- Implement automated plot-generating scripts comparing cosmos to different diffsky models
+    - https://github.com/ArgonneCPAC/diffsky/pull/291
+- Improve bookkeeping in cross-matching jobs
+    - https://github.com/ArgonneCPAC/diffsky/pull/293
+    - https://github.com/ArgonneCPAC/diffsky/pull/294
+    - https://github.com/ArgonneCPAC/diffsky/pull/299
+    - https://github.com/ArgonneCPAC/diffsky/pull/300
+- Compute photometry in batches when making mocks
+    - https://github.com/ArgonneCPAC/diffsky/pull/295
+- Adopt astropy units in mock metadata
+    - https://github.com/ArgonneCPAC/diffsky/pull/303
+- Add peculiar velocities to output mock
+    - https://github.com/ArgonneCPAC/diffsky/pull/304
+
+
 0.3.2 (2025-12-04)
 -------------------
 - Resolve bug in the treatment of scatter in the SSP errors.

{diffsky-0.3.2 → diffsky-0.3.3}/PKG-INFO

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

diffsky-0.3.3/diffsky/_version.py

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

{diffsky-0.3.2 → diffsky-0.3.3}/diffsky/burstpop/diagnostics/plot_fburstpop.py

@@ -1,10 +1,17 @@
 """ """
 
 import numpy as np
-from matplotlib import pyplot as plt
 
 from ..fburstpop import DEFAULT_FBURSTPOP_PARAMS, get_lgfburst_from_fburstpop_params
 
+try:
+    from matplotlib import pyplot as plt
+
+    HAS_MATPLOTLIB = True
+except ImportError:
+    HAS_MATPLOTLIB = False
+MATPLOTLIB_MSG = "Must have matplotlib installed to use this function"
+
 
 def make_fburstpop_comparison_plot(
     params,
@@ -28,6 +35,8 @@ def make_fburstpop_comparison_plot(
         filename of the output figure
 
     """
+    assert HAS_MATPLOTLIB, MATPLOTLIB_MSG
+
     nsm, nsfr = 250, 250
     logsm_grid = np.linspace(7, 12, nsm)
     logssfr_grid = np.linspace(-13, -8, nsfr)

{diffsky-0.3.2 → diffsky-0.3.3}/diffsky/burstpop/diagnostics/plot_tburstpop.py

@@ -2,11 +2,18 @@
 
 import numpy as np
 from dsps.sfh import diffburst
-from matplotlib import lines as mlines
-from matplotlib import pyplot as plt
 
 from .. import tburstpop as tbp
 
+try:
+    from matplotlib import lines as mlines
+    from matplotlib import pyplot as plt
+
+    HAS_MATPLOTLIB = True
+except ImportError:
+    HAS_MATPLOTLIB = False
+MATPLOTLIB_MSG = "Must have matplotlib installed to use this function"
+
 
 def make_tburstpop_comparison_plot(
     params,
@@ -30,6 +37,7 @@ def make_tburstpop_comparison_plot(
         filename of the output figure
 
     """
+    assert HAS_MATPLOTLIB, MATPLOTLIB_MSG
     lgyrarr = np.linspace(5, 9.05, 100)
 
     logsmarr = np.array((9.0, 9.0, 12.0, 12.0))

{diffsky-0.3.2 → diffsky-0.3.3}/diffsky/burstpop/diagnostics/tests/test_plot_fburstpop.py

@@ -2,11 +2,19 @@
 
 import os
 
-from ..plot_fburstpop import DEFAULT_FBURSTPOP_PARAMS, make_fburstpop_comparison_plot
+import pytest
+
+from ..plot_fburstpop import (
+    DEFAULT_FBURSTPOP_PARAMS,
+    HAS_MATPLOTLIB,
+    MATPLOTLIB_MSG,
+    make_fburstpop_comparison_plot,
+)
 
 _THIS_DRNAME = os.path.dirname(os.path.abspath(__file__))
 
 
+@pytest.mark.skipif(not HAS_MATPLOTLIB, reason=MATPLOTLIB_MSG)
 def test_make_freqburst_comparison_plot():
     fn = os.path.join(_THIS_DRNAME, "dummy.png")
     make_fburstpop_comparison_plot(DEFAULT_FBURSTPOP_PARAMS, fname=fn)

{diffsky-0.3.2 → diffsky-0.3.3}/diffsky/burstpop/diagnostics/tests/test_plot_tburstpop.py

@@ -2,12 +2,19 @@
 
 import os
 
+import pytest
+
 from ... import tburstpop as tbp
-from ..plot_tburstpop import make_tburstpop_comparison_plot
+from ..plot_tburstpop import (
+    HAS_MATPLOTLIB,
+    MATPLOTLIB_MSG,
+    make_tburstpop_comparison_plot,
+)
 
 _THIS_DRNAME = os.path.dirname(os.path.abspath(__file__))
 
 
+@pytest.mark.skipif(not HAS_MATPLOTLIB, reason=MATPLOTLIB_MSG)
 def test_make_freqburst_comparison_plot():
     fn = os.path.join(_THIS_DRNAME, "dummy.png")
     make_tburstpop_comparison_plot(tbp.DEFAULT_TBURSTPOP_PARAMS, fname=fn)

{diffsky-0.3.2 → diffsky-0.3.3}/diffsky/data_loaders/cosmos20_loader.py

@@ -13,6 +13,15 @@ except ImportError:
     HAS_ASTROPY = False
 
 
+Z_MIN, Z_MAX = 0.4, 2.5
+
+HSC_MAG_NAMES = ["HSC_g_MAG", "HSC_r_MAG", "HSC_i_MAG", "HSC_z_MAG", "HSC_y_MAG"]
+UVISTA_MAG_NAMES = ["UVISTA_Y_MAG", "UVISTA_J_MAG", "UVISTA_H_MAG", "UVISTA_Ks_MAG"]
+COSMOS_TARGET_MAGS = [*HSC_MAG_NAMES, *UVISTA_MAG_NAMES]
+MAGI_THRESH = 25.5
+
+NANFILL = -999.0
+
 COSMOS20_BASENAME = "COSMOS2020_Farmer_processed_hlin.fits"
 
 SKY_AREA = 1.21  # square degrees
@@ -103,4 +112,85 @@ def load_cosmos20(
     else:
         return cat
 
-    return cat
+
+def apply_nan_cuts(cosmos, mag_names=COSMOS_TARGET_MAGS):
+    """Remove any galaxy with a NaN in any column storing a target magnitude
+
+    Parameters
+    ----------
+    cosmos : astropy Table
+
+    mag_names : list of strings
+
+    Returns
+    -------
+    cosmos : astropy Table
+        Catalog after applying a NaN cut on `photoz` and any colname in mag_names
+
+    """
+    msk_has_nan = np.isnan(cosmos["photoz"])
+    for name in mag_names:
+        x = np.nan_to_num(
+            cosmos[name], copy=True, nan=NANFILL, posinf=NANFILL, neginf=NANFILL
+        )
+        msk_has_nan = msk_has_nan | (x == NANFILL)
+
+    cosmos = cosmos[~msk_has_nan]
+    return cosmos
+
+
+def get_is_complete_mask(cosmos, z_min=Z_MIN, z_max=Z_MAX, magi_thresh=MAGI_THRESH):
+    """Compute mask to define the redshift and i-mag threshold for our target data
+
+    Parameters
+    ----------
+    cosmos : astropy Table
+
+    z_min, z_max : float
+        Galaxies outside this range will be excluded
+
+    magi_thresh : float
+        Galaxies fainter than this apparent magnitude will be excluded
+
+    Returns
+    -------
+    msk_is_complete : array, dtype bool
+        Boolean mask defining galaxies that pass the completeness cut
+
+    """
+    msk_redshift = (cosmos["photoz"] > z_min) & (cosmos["photoz"] < z_max)
+    msk_i_thresh = cosmos["HSC_i_MAG"] < magi_thresh
+    msk_is_complete = msk_redshift & msk_i_thresh
+    return msk_is_complete
+
+
+def get_color_outlier_mask(cosmos, mag_names, p_cut=0.5):
+    """Compute mask to define extreme outliers in color space
+
+    Parameters
+    ----------
+    cosmos : astropy Table
+
+    mag_names : list of strings
+        Column names defining the colors for which outliers will be excluded
+
+    p_cut : float, optional
+        Value in the range [0, 100] defining a percentile cut
+
+    Returns
+    -------
+    msk_is_not_outlier : array, dtype bool
+        Boolean mask defining galaxies that pass the outlier cut
+
+    """
+    msk_is_outlier = np.zeros(len(cosmos)).astype(bool)
+    for name0, name1 in zip(mag_names[0:], mag_names[1:]):
+        c0 = cosmos[name0]
+        c1 = cosmos[name1]
+        color = c0 - c1
+        lo, hi = np.percentile(color, (p_cut, 100.0 - p_cut))
+        msk_is_outlier = msk_is_outlier | (color < lo) | (color > hi)
+
+    msk_is_not_outlier = ~msk_is_outlier
+
+    return msk_is_not_outlier

{diffsky-0.3.2 → diffsky-0.3.3}/diffsky/data_loaders/hacc_utils/data_validation/validate_lc_mock.py

@@ -9,10 +9,10 @@ from diffstar import DEFAULT_DIFFSTAR_PARAMS
 from dsps.cosmology.flat_wcdm import age_at_z
 
 from .... import phot_utils
-from ....experimental import dbk_from_mock2
+from ....experimental import dbk_phot_from_mock
 from ....experimental import precompute_ssp_phot as psspp
 from ....param_utils import diffsky_param_wrapper as dpw
-from .. import lc_mock_repro as lcmp
+from .. import lc_mock as lcmp
 from .. import load_flat_hdf5, load_lc_cf
 
 REQUIRED_METADATA_ATTRS = ("creation_date", "README", "mock_version_name")
@@ -93,17 +93,38 @@ def check_all_columns_are_finite(fn_lc_mock, data=None):
 
 def check_all_data_columns_have_metadata(fn_lc_mock):
 
+    try:
+        from astropy import units as u
+        from astropy.cosmology import units as cu
+
+        u.add_enabled_units(cu)
+
+    except ImportError:
+        raise ImportError("Must have astropy installed to attach units to metadata")
+
     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
+                unit_string = hdf["data/" + key].attrs["unit"]
+                u.Unit(unit_string)
+
+            except (KeyError,):
+                s = f"{key} is missing metadata unit"
+                msg.append(s)
+            except (AssertionError, ValueError):
+                s = f"Units not recognized in metadata for `{key}` column"
+                msg.append(s)
+
+            try:
                 description = hdf["data/" + key].attrs["description"]
                 assert len(description) > 0
-            except (KeyError, AssertionError):
-                s = f"{key} is missing metadata"
+            except (KeyError,):
+                s = f"{key} is missing metadata description"
+                msg.append(s)
+            except (AssertionError, ValueError):
+                s = f"No metadata description of `{key}` column"
                 msg.append(s)
     return msg
 
@@ -322,7 +343,7 @@ def check_recomputed_photometry(fn_lc_mock, n_test=200, return_results=False):
         sim_info.cosmo_params,
         sim_info.fb,
     )
-    _res = dbk_from_mock2._reproduce_mock_dbk_kern(*args)
+    _res = dbk_phot_from_mock._reproduce_mock_dbk_kern(*args)
     (
         phot_info,
         phot_randoms,

diffsky-0.3.2/diffsky/data_loaders/hacc_utils/lc_mock_repro.py → diffsky-0.3.3/diffsky/data_loaders/hacc_utils/lc_mock.py

@@ -37,6 +37,7 @@ from ...experimental.kernels import mc_phot_kernels as mcpk
 from ...experimental.size_modeling import disk_bulge_sizes as dbs
 from ...fake_sats import halo_boundary_functions as hbf
 from ...fake_sats import nfw_config_space as nfwcs
+from ...fake_sats import vector_utilities as vecu
 from .. import load_flat_hdf5
 from . import lightcone_utils as hlu
 from . import load_lc_cf
@@ -52,18 +53,20 @@ 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_NFW_KEYS_OUT = (
+    "x_nfw",
+    "y_nfw",
+    "z_nfw",
+    "ra_nfw",
+    "dec_nfw",
+)
 LC_DATA_KEYS_OUT = (
     "core_tag",
     "x",
     "y",
     "z",
-    "x_nfw",
-    "y_nfw",
-    "z_nfw",
     "top_host_idx",
     "central",
-    "ra_nfw",
-    "dec_nfw",
     "redshift_true",
     "stepnum",
     "lc_patch",
@@ -85,14 +88,18 @@ _ORIEN_PATS = (
 ORIENTATION_KEYS = [pat.format("disk") for pat in _ORIEN_PATS]
 ORIENTATION_KEYS.extend([pat.format("bulge") for pat in _ORIEN_PATS])
 
-DIFFSKY_DATA_KEYS_OUT = (
+DIFFSKY_DATA_NFW_HOST_KEYS_OUT = (
     "x_host",
     "y_host",
     "z_host",
+    "logmp_obs_host",
+)
+DIFFSKY_DATA_KEYS_OUT = (
     "vx",
     "vy",
     "vz",
-    "logmp_obs_host",
+    "vpec",
+    "msk_v0",
     "has_diffmah_fit",
     "logmp0",
     "logmp_obs",
@@ -264,6 +271,52 @@ def write_lc_sfh_mock_to_disk(fnout, lc_data, diffsky_data):
             hdf_out[key_out] = diffsky_data[key]
 
 
+def get_imputed_velocity(vx, vy, vz, ran_key, std_v=500.0):
+    """Overwrite zero-valued velocities with random normal data"""
+    msk_imputed = (vx == 0) & (vy == 0) & (vz == 0)
+    vran = jran.normal(ran_key, (vx.size, 3)) * std_v
+    v = np.array((vx, vy, vz)).T
+    imputed_v = np.where(msk_imputed.reshape((-1, 1)), vran, v)
+    vx_imputed = imputed_v[:, 0]
+    vy_imputed = imputed_v[:, 1]
+    vz_imputed = imputed_v[:, 2]
+    return vx_imputed, vy_imputed, vz_imputed, msk_imputed
+
+
+def write_batched_lc_sfh_mock_to_disk(fnout, lc_data, diffsky_data):
+    """"""
+    write_batched_mock_data(fnout, lc_data, LC_DATA_KEYS_OUT, dataset="data")
+    write_batched_mock_data(fnout, diffsky_data, DIFFSKY_DATA_KEYS_OUT, dataset="data")
+
+    ra, dec = hlu._get_lon_lat_from_theta_phi(lc_data["theta"], lc_data["phi"])
+    ra_dec_dict = dict(ra=ra, dec=dec)
+    write_batched_mock_data(
+        fnout, ra_dec_dict, list(ra_dec_dict.keys()), dataset="data"
+    )
+
+
+def write_batched_lc_sed_mock_to_disk(
+    fnout, phot_info, lc_data, diffsky_data, filter_nicknames
+):
+    write_batched_lc_sfh_mock_to_disk(fnout, lc_data, diffsky_data)
+
+    phot_dict = dict()
+    for iband, name in enumerate(filter_nicknames):
+        phot_dict[name] = phot_info["obs_mags"][:, iband]
+    write_batched_mock_data(fnout, phot_dict, list(phot_dict.keys()), dataset="data")
+
+    phot_info_colnames = [
+        *DEFAULT_BURST_PARAMS._fields,
+        *DEFAULT_DUST_PARAMS._fields,
+        "mc_sfh_type",
+        *PHOT_INFO_KEYS_OUT,
+    ]
+    write_batched_mock_data(fnout, phot_info, phot_info_colnames, dataset="data")
+
+    diffsky_data_colnames = [*MORPH_KEYS_OUT, *BLACK_HOLE_KEYS_OUT]
+    write_batched_mock_data(fnout, diffsky_data, diffsky_data_colnames, dataset="data")
+
+
 def write_lc_sed_mock_to_disk(
     fnout, phot_info, lc_data, diffsky_data, filter_nicknames
 ):
@@ -295,7 +348,6 @@ def write_lc_dbk_sed_mock_to_disk(
     fnout, phot_info, lc_data, diffsky_data, filter_nicknames
 ):
     write_lc_sed_mock_to_disk(fnout, phot_info, lc_data, diffsky_data, filter_nicknames)
-
     with h5py.File(fnout, "a") as hdf_out:
         for iband, name in enumerate(filter_nicknames):
             hdf_out["data"][name + "_bulge"] = phot_info["obs_mags_bulge"][:, iband]
@@ -304,6 +356,52 @@ def write_lc_dbk_sed_mock_to_disk(
         hdf_out["data"]["fknot"] = phot_info["fknot"]
 
 
+def write_batched_lc_dbk_sed_mock_to_disk(
+    fnout, phot_info, lc_data, diffsky_data, filter_nicknames
+):
+    write_batched_lc_sed_mock_to_disk(
+        fnout, phot_info, lc_data, diffsky_data, filter_nicknames
+    )
+    dbk_phot_dict = dict()
+    for iband, name in enumerate(filter_nicknames):
+        dbk_phot_dict[name + "_bulge"] = phot_info["obs_mags_bulge"][:, iband]
+        dbk_phot_dict[name + "_disk"] = phot_info["obs_mags_disk"][:, iband]
+        dbk_phot_dict[name + "_knots"] = phot_info["obs_mags_knots"][:, iband]
+        dbk_phot_dict["fknot"] = phot_info["fknot"]
+    write_batched_mock_data(
+        fnout, dbk_phot_dict, list(dbk_phot_dict.keys()), dataset="data"
+    )
+
+
+def add_peculiar_velocity_to_mock(
+    lc_data, diffsky_data, ran_key=None, impute_vzero=True
+):
+    # Patch v==0 galaxies
+    if impute_vzero:
+        assert ran_key is not None, "Must pass ran_key when impute_vzero=True"
+
+        vx, vy, vz, msk_imputed = get_imputed_velocity(
+            diffsky_data["vx"], diffsky_data["vy"], diffsky_data["vz"], ran_key
+        )
+    else:
+        vx = diffsky_data["vx"]
+        vy = diffsky_data["vy"]
+        vz = diffsky_data["vz"]
+        msk_imputed = np.zeros(len(vx)).astype(bool)
+
+    diffsky_data["vx"] = vx
+    diffsky_data["vy"] = vy
+    diffsky_data["vz"] = vz
+    diffsky_data["msk_v0"] = msk_imputed
+
+    X = np.array((lc_data["x"], lc_data["y"], lc_data["z"])).T
+    V = np.array((diffsky_data["vx"], diffsky_data["vy"], diffsky_data["vz"])).T
+    Xnorm = vecu.normalized_vectors(X)
+    diffsky_data["vpec"] = vecu.elementwise_dot(V, Xnorm)
+
+    return diffsky_data
+
+
 def add_diffmah_properties_to_mock(diffsky_data, redshift_true, sim_info, ran_key):
     diffsky_data["t_obs"] = flat_wcdm.age_at_z(redshift_true, *sim_info.cosmo_params)
 
@@ -535,3 +633,39 @@ def concatenate_batched_phot_data(phot_batches):
             raise ValueError(f"Unable to concatenate diffsky_data['{key}']")
 
     return phot_info, lc_data, diffsky_data
+
+
+def write_batched_mock_data(fn_out, data_batch, colnames_out, dataset="data"):
+
+    with h5py.File(fn_out, "a") as hdf_out:
+        if dataset not in hdf_out:
+            grp = hdf_out.create_group(dataset)
+        else:
+            grp = hdf_out[dataset]
+
+        for key in colnames_out:
+            # Ensure data_batch[key] is an array
+            arr = np.asarray(data_batch[key])
+            if arr.ndim == 0:
+                arr = np.array([arr])
+
+            if key not in grp:
+                maxshape = (None,) + arr.shape[1:]
+
+                # Explicitly set chunk size instead of chunks=True
+                chunk_size = min(1000, arr.shape[0])
+                chunks = (chunk_size,) + arr.shape[1:]
+
+                grp.create_dataset(key, data=arr, maxshape=maxshape, chunks=chunks)
+            else:
+                current_size = grp[key].shape[0]
+                new_size = current_size + arr.shape[0]
+                try:
+                    grp[key].resize(new_size, axis=0)
+                    grp[key][current_size:new_size] = arr
+                except (TypeError, RuntimeError) as e:
+                    raise TypeError(
+                        f"Tried to resize column `{key}`. "
+                        f"Original error: {e}. "
+                        f"Dataset chunks: {grp[key].chunks}"
+                    )