Feature/159 psf output from trained model

changelog.d/20260123_160331_jennifer.pollack_159_psf_output_from_trained_model.md

@@ -0,0 +1,45 @@
+<!--
+A new scriv changelog fragment.
+
+Uncomment the section that is right (remove the HTML comment wrapper).
+For top level release notes, leave all the headers commented out.
+-->
+
+<!--
+### Breaking changes
+
+- A bullet item for the Breaking changes category.
+
+-->
+
+### New features
+
+- Added PSF inference capabilities for generating broadband (polychromatic) PSFs from trained models given star positions and SEDs
+- Introduced `PSFInferenceEngine` class to centralize training, simulation, metrics, and inference workflows
+- Added `run_type` attribute to `DataHandler` supporting training, simulation, metrics, and inference modes
+- Implemented `ZernikeInputsFactory` class for building `ZernikeInputs` instances based on run type
+- Added `psf_model_loader.py` module for centralized model weights loading
+
+
+### Bug fixes
+
+- Fix logger formatting for relative RMSE metrics in `metrics.py` (values were not being displayed)
+
+<!--
+### Performance improvements
+
+- A bullet item for the Performance improvements category.
+
+-->
+
+### Internal changes
+
+- Refactored `TFPhysicalPolychromatic` and related modules to separate training vs. inference behavior
+- Enhanced `ZernikeInputs` data class with intelligent assembly based on run type and available data
+- Implemented hybrid loading pattern with eager loading in constructors and lazy-loading via property decorators
+- Centralized PSF data extraction in `data_handler` module
+- Improved code organization with new `tf_utils.py` module in `psf_models` sub-package
+- Updated configuration handling to support inference workflows via `inference_config.yaml`
+- Fixed incorrect argument name in `DataHandler` that prevented proper TensorFlow data type conversion
+- Removed deprecated `get_obs_positions` method
+- Updated documentation to include inference package

changelog.d/20260210_150506_jennifer.pollack_159_psf_output_from_trained_model.md

@@ -0,0 +1,37 @@
+<!--
+A new scriv changelog fragment.
+
+Uncomment the section that is right (remove the HTML comment wrapper).
+For top level release notes, leave all the headers commented out.
+-->
+
+<!--
+### Breaking changes
+
+- A bullet item for the Breaking changes category.
+
+-->
+<!--
+### New features
+
+- A bullet item for the New features category.
+
+-->
+<!--
+### Bug fixes
+
+- A bullet item for the Bug fixes category.
+
+-->
+<!--
+### Performance improvements
+
+- A bullet item for the Performance improvements category.
+
+-->
+
+### Internal changes
+
+- Remove deprecated/optional import tensorflow-addons statement from tf_layers.py
+
+

config/inference_conf.yaml

@@ -0,0 +1,37 @@
+
+inference:
+  # Inference batch size
+  batch_size: 16
+  # Cycle to use for inference. Can be: 1, 2, ...
+  cycle: 2
+
+  # Paths to the configuration files and trained model directory
+  configs:
+    # Path to the directory containing the trained model
+    trained_model_path: /path/to/trained/model/
+
+    # Subdirectory name of the trained model, e.g. psf_model
+    model_subdir: model
+
+    # Relative Path to the training configuration file used to train the model
+    trained_model_config_path: config/training_config.yaml
+
+    # Path to the data config file (this could contain prior information)
+    data_config_path: 
+
+  # The following parameters will overwrite the `model_params` in the training config file.
+  model_params:
+    # Num of wavelength bins to reconstruct polychromatic objects.
+    n_bins_lda: 8
+
+    # Downsampling rate to match the oversampled model to the specified telescope's sampling.
+    output_Q: 1
+
+    # Dimension of the pixel PSF postage stamp
+    output_dim: 64
+
+    # Flag to perform centroid error correction
+    correct_centroids: False
+
+    # Flag to perform CCD misalignment error correction
+    add_ccd_misalignments: True

docs/source/api.rst

@@ -8,6 +8,7 @@ This section contains the API reference for the main packages in WaveDiff.
    :recursive:
 
    wf_psf.data
+   wf_psf.inference
    wf_psf.metrics
    wf_psf.plotting
    wf_psf.psf_models

docs/source/conf.py

@@ -21,7 +21,7 @@
 else:
     copyright = f"{start_year}, CosmoStat"
 author = "CosmoStat"
-release = "3.0.0"
+release = "3.1.0"
 
 # -- General configuration ---------------------------------------------------
 # https://www.sphinx-doc.org/en/master/usage/configuration.html#general-configuration

docs/source/configuration.md

@@ -137,13 +137,13 @@ model_params:
   reference_shifts: [-1/3, -1/3]   # Euclid-like default shifts
 
   # Obscuration / geometry
-  obscuration_rotation_angle: 0    # Degrees (multiple of 90); counterclockwise rotation. 
+  obscuration_rotation_angle: 0    # Degrees (multiple of 90); counterclockwise rotation.
 
   # CCD misalignments input file path
   ccd_misalignments_input_path: /path/to/ccd_misalignments_file.txt
-    
+
   # Boolean to use sample weights based on the noise standard deviation estimation
-  use_sample_weights: True 
+  use_sample_weights: True
 
   # Sample weight generalised sigmoid function
   sample_weights_sigmoid:
@@ -220,7 +220,6 @@ training_hparams:
     n_epochs_non_params: [100, 120]
 ```
 
-
 (metrics_config)=
 ## `metrics_config.yaml`  — Metrics Configuration
 
@@ -402,7 +401,10 @@ plotting_params:
 ### 4. Example Directory Structure
 Below is an example of three WaveDiff runs stored under a single parent directory:
 
-```
+**Example Directory Structure**
+Below is an example of three WaveDiff runs stored under a single parent directory:
+
+```arduino
 wf-outputs/
 ├── wf-outputs-202305271829
 │   ├── config

pyproject.toml

@@ -11,7 +11,7 @@ maintainers = [
 
 description = 'A software framework to perform Differentiable wavefront-based PSF modelling.'
 dependencies = [
-    "numpy>=1.26.4,<2.0",
+    "numpy>=1.18,<1.24",
     "scipy",
     "tensorflow==2.11.0",
     "tensorflow-estimator",
@@ -24,7 +24,7 @@ dependencies = [
     "seaborn",
 ]
 
-version = "3.0.0"
+version = "3.1.0"
 
 [project.optional-dependencies]
 docs = [

src/wf_psf/__init__.py

@@ -2,6 +2,6 @@
 
 # Dynamically import modules to trigger side effects when wf_psf is imported
 importlib.import_module("wf_psf.psf_models.psf_models")
-importlib.import_module("wf_psf.psf_models.psf_model_semiparametric")
-importlib.import_module("wf_psf.psf_models.psf_model_physical_polychromatic")
-importlib.import_module("wf_psf.psf_models.tf_psf_field")
+importlib.import_module("wf_psf.psf_models.models.psf_model_semiparametric")
+importlib.import_module("wf_psf.psf_models.models.psf_model_physical_polychromatic")
+importlib.import_module("wf_psf.psf_models.tf_modules.tf_psf_field")

src/wf_psf/utils/centroids.py → src/wf_psf/data/centroids.py

@@ -8,10 +8,80 @@
 
 import numpy as np
 import scipy.signal as scisig
-from wf_psf.utils.preprocessing import shift_x_y_to_zk1_2_wavediff
+from fractions import Fraction
 from typing import Optional
 
 
+def compute_centroid_correction(
+    model_params, centroid_dataset, batch_size: int = 1
+) -> np.ndarray:
+    """Compute centroid corrections using Zernike polynomials.
+
+    This function calculates the Zernike contributions required to match the centroid
+    of the WaveDiff PSF model to the observed star centroids, processing in batches.
+
+    Parameters
+    ----------
+    model_params : RecursiveNamespace
+        An object containing parameters for the PSF model, including pixel sampling and initial centroid window parameters.
+
+    centroid_dataset : dict
+        Dictionary containing star data needed for centroiding:
+        - "stamps" : np.ndarray
+            Array of star postage stamps (required).
+        - "masks" : Optional[np.ndarray]
+            Array of star masks (optional, can be None).
+
+    batch_size : int, optional
+        The batch size to use when processing the stars. Default is 16.
+
+    Returns
+    -------
+    zernike_centroid_array : np.ndarray
+         A 2D NumPy array of shape `(n_stars, 3)`, where `n_stars` is the number of
+        observed stars. The array contains the computed Zernike (Z1, Z2) contributions,
+        with zero padding applied to the first column to ensure a consistent shape.
+    """
+    # Retrieve stamps and masks from centroid_dataset
+    star_postage_stamps = centroid_dataset.get("stamps")
+    star_masks = centroid_dataset.get("masks")  # may be None
+
+    if star_postage_stamps is None:
+        raise ValueError("centroid_dataset must contain 'stamps'")
+
+    pix_sampling = model_params.pix_sampling * 1e-6  # Change units from [um] to [m]
+
+    reference_shifts = [
+        float(Fraction(value)) for value in model_params.reference_shifts
+    ]
+
+    n_stars = len(star_postage_stamps)
+    zernike_centroid_array = []
+
+    # Batch process the stars
+    for i in range(0, n_stars, batch_size):
+        batch_postage_stamps = star_postage_stamps[i : i + batch_size]
+        batch_masks = star_masks[i : i + batch_size] if star_masks is not None else None
+
+        # Compute Zernike 1 and Zernike 2 for the batch
+        zk1_2_batch = -1.0 * compute_zernike_tip_tilt(
+            batch_postage_stamps, batch_masks, pix_sampling, reference_shifts
+        )
+
+        # Zero pad array for each batch and append
+        zernike_centroid_array.append(
+            np.pad(
+                zk1_2_batch,
+                pad_width=[(0, 0), (1, 0)],
+                mode="constant",
+                constant_values=0,
+            )
+        )
+
+    # Combine all batches into a single array
+    return np.concatenate(zernike_centroid_array, axis=0)
+
+
 def compute_zernike_tip_tilt(
     star_images: np.ndarray,
     star_masks: Optional[np.ndarray] = None,
@@ -58,6 +128,8 @@ def compute_zernike_tip_tilt(
     - This function processes all images at once using vectorized operations.
     - The Zernike coefficients are computed in the WaveDiff convention.
     """
+    from wf_psf.data.data_zernike_utils import shift_x_y_to_zk1_2_wavediff
+
     # Vectorize the centroid computation
     centroid_estimator = CentroidEstimator(
         im=star_images, mask=star_masks, sigma_init=sigma_init, n_iter=n_iter
@@ -178,6 +250,18 @@ def __init__(
         self, im, mask=None, sigma_init=7.5, n_iter=5, auto_run=True, xc=None, yc=None
     ):
         """Initialize class attributes."""
+        # Convert to np.ndarray if not already
+        im = np.asarray(im)
+        if mask is not None:
+            mask = np.asarray(mask)
+
+        # Check im dimensions convert to batch, if 2D
+        if im.ndim == 2:
+            # Single stamp → convert to batch of one
+            im = np.expand_dims(im, axis=0)
+        elif im.ndim != 3:
+            raise ValueError(f"Expected 2D or 3D input, got shape {im.shape}")
+
         self.im = im
         self.mask = mask
         if self.mask is not None: