Skip to content

Feature/159 psf output from trained model#163

Open
jeipollack wants to merge 129 commits intodevelopfrom
feature/159-psf-output-from-trained-model
Open

Feature/159 psf output from trained model#163
jeipollack wants to merge 129 commits intodevelopfrom
feature/159-psf-output-from-trained-model

Conversation

@jeipollack
Copy link
Contributor

@jeipollack jeipollack commented May 13, 2025 •
edited
Loading

This PR adds inference capabilities to the WaveDiff / WF-PSF framework, allowing the user to generate broadband (polychromatic) PSFs from a trained model, given a list of positions and SEDs.

It had a refactor PR merged into it in order to carry on with the development which is why it contains >100 commits (unfortunately and not a practice that will be repeated).

Closes #159 (Validation runs described there)

Key Changes

  • Introduced psf_inference.py and associated test_inference packages for PSF inference workflows
  • Refactored TFPhysicalPolychromatic and related modules to separate training vs. inference behaviour
  • Added PSFInferenceEngine class to centralise training, simulation, metrics, and inference logic
  • Moved model weights loading to psf_model_loader.py
  • Centralised PSF data extraction in data_handler
  • Updated configuration handling for inference (inference_config.yaml)
  • Updated documentation (small) to include inference package
  • Core Refactoring from PR 165 refactor tfphysicalpolychromatic to cleanly separate training vs inference behaviour #166:
    • New run_type attribute: Added to data_handler supporting training, simulation, metrics, and inference modes
    • ZernikeInputsFactory class: New factory class that builds ZernikeInputs instances based on the specified run type
    • Enhanced ZernikeInputs data class: Now intelligently assembles Zernike contribution lists based on run_type and available data (including priors and source positions for CCD misalignment corrections)
    • Hybrid loading pattern: Refactored to use eager loading in constructors combined with lazy-loading via property decorators, reducing constructor overhead and improving modularity
    • Code organization: Improved module structure including new tf_utils.py module in the psf_models sub-package
    • Bug fix: Fixed incorrect argument name in DataHandler class that prevented training dataset images from being converted to TensorFlow data type as expected
    • Helper Methods: Added several helper methods to assist with preparing datasets and ensure data type conversion is performed as expected
    • Deprecations: Removed get_obs_positions as it was no longer needed
  • Bug fix: Fix logger formatting for relative RMSE metrics in metrics.py (values were not being displayed)
  • Changelog Fragment added describing updates

Testing

  • CI tests pass successfully
  • Unit tests cover:
    • Loading trained models and other aspects of inference setup
    • Generating PSFs for multiple positions/SEDs and error handling
    • Consistency of outputs and tensor handling
    • Clearing cached values
  • Verified backward compatibility with existing develop code
  • Confirmed inferred PSFs from test datasets match expected results with high accuracy

@jeipollack jeipollack self-assigned this May 13, 2025
@jeipollack jeipollack added the enhancement New feature or request label May 13, 2025
@jeipollack jeipollack changed the title Feature/159 psf output from trained model Draft:Feature/159 psf output from trained model May 13, 2025
@jeipollack jeipollack mentioned this pull request Jun 13, 2025
Closed
@jeipollack jeipollack changed the base branch from develop to case_study_psf_decontamination July 3, 2025 07:21
Base automatically changed from case_study_psf_decontamination to develop August 5, 2025 10:32
@jeipollack jeipollack force-pushed the feature/159-psf-output-from-trained-model branch from f15c0f2 to 89da2e0 Compare August 5, 2025 13:39
@jeipollack jeipollack mentioned this pull request Sep 1, 2025
Open
@jeipollack jeipollack changed the title Draft:Feature/159 psf output from trained model Feature/159 psf output from trained model Nov 12, 2025
@jeipollack jeipollack requested a review from sfarrens November 12, 2025 10:39
@jeipollack jeipollack linked an issue Nov 19, 2025 that may be closed by this pull request
Add function to output PSF for shape measurement given a trained model #159
Open
@jeipollack jeipollack added the documentation Improvements or additions to documentation label Nov 20, 2025
@jeipollack jeipollack changed the title Feature/159 psf output from trained model Draft: Feature/159 psf output from trained model Nov 20, 2025
@jeipollack jeipollack force-pushed the feature/159-psf-output-from-trained-model branch 2 times, most recently from c199e05 to 9f64df3 Compare December 8, 2025 12:21
@jeipollack jeipollack force-pushed the feature/159-psf-output-from-trained-model branch from e995e7c to 6b2a1fe Compare January 22, 2026 15:30
@jeipollack jeipollack changed the title Draft: Feature/159 psf output from trained model Feature/159 psf output from trained model Jan 26, 2026
@jeipollack
Copy link
Contributor Author

jeipollack commented Jan 26, 2026

it would be very ideal to squash and merge this commit in order to simplify rebasing the other branches.

@jeipollack jeipollack removed the request for review from sfarrens January 26, 2026 12:12
@jeipollack jeipollack changed the title Feature/159 psf output from trained model DRAFT: Feature/159 psf output from trained model Jan 26, 2026
@jeipollack jeipollack changed the title DRAFT: Feature/159 psf output from trained model Feature/159 psf output from trained model Jan 26, 2026
Jennifer Pollack added 6 commits February 2, 2026 16:32
Correct doc string, format errors, unused imports, type hints, etc
0884737
Update documentation for v3.0.0: replace modules.rst with api.rst, cl…
532e985 
…ean up submodule docs, update conf.py, and update pyproject.toml version
Update API reference and clean up package/module docstrings
be61b0f 
- Removed old auto-generated wf_psf.rst from _autosummary
- Updated toc.rst (fixed tab issues, reference api.rst)
- Added api.rst with :recursive: directive and wf_psf.run entrypoint
- Refined __init__.py docstrings for all subpackages for clarity and consistency
- Updated module-level docstrings (purpose, authors, TensorFlow notes, etc.)
Remove auto-generated _autosummary from repo; will be built in CD
afa2923
Add arduino formatting to directory structure example
6ecb9ba
Revise WaveDiff how-to guide in doc (part 1)
11bad16 
- Shorten sections to make developer-friendly
- Added new structure for each section: Purppose, Key Fields, Notes, General Notes, etc, where applicable
- Partial completion - new PR is required to verify optional versus required settings
Jennifer Pollack and others added 28 commits February 2, 2026 16:46
Replace np.array with Tensorflow tensors in unit test and fixtures, r…
829136a 
…eformat files
@jeipollack
Eagerly initialise trainable layers in physical poly model constructo…
96460a5 
…r required for evaluation/inference
@jeipollack
fix: use expect_partial() when loading model weights for evaluation
06d715b 
- Add status handling to model.load_weights() call
- Use expect_partial() to suppress warnings about unused optimizer state
- Allows successful weight loading for metrics evaluation when checkpoint contains training artifacts
@jeipollack
Add memory cleanup after training completion
5b3dee3 
- Delete model reference and run garbage collection
- Clear TensorFlow session to free GPU memory
- Prevents OOM issues in subsequent operations or multiple training runs
@jeipollack
refactor: centralise PSF data extraction in data_handler
6a2d24d 
- Introduce unified `get_data_array` for training/metrics/inference access
- Add helpers `extract_star_data` and `_get_inference_data`
- Remove redundant `get_np_obs_positions`
- Move data handling logic out of `compute_centroid_corrections`
- Standardise `centroid_dataset` as dict (stamps + optional masks)
- Support optional keys (masks, priors) via `allow_missing`
- Improve and unify docstrings for data extraction utilities
- Add optional "sources" and "masks" attributes to `PSFInference
- Add `correct_centroids` and `add_ccd_misalignments` as options to
  inference_config.yaml
@jeipollack
Add and options to inference_config.yaml (forgot to stage with previo…
8376fe1 
…us commit)
@jeipollack
Update PSFInference doc string with new optional attributes
77434be
@jeipollack
Rename _get_inference_data to _get_direct_data
76a449c
Reformat with black
b0874e7
@jeipollack
Correct type hint errors
0e85c0f
@jeipollack
Remove unused import
2bdd622
@jeipollack
Replace call to deprecated get_np_obs_positions with get_data_array
6134d8e
Remove unused imports and reformat
e309487
Update fixtures and unit tests
8832584
Reformat with black
529fcd3
Remove outdated back module
6c3e9d0
Remove unused import left after rebase
af12a6c
Update doc strings and cache handling
07e88c7 
- Updated classes and methods with complete doc strings
- Added two cache clearing methods to PSFInference and PSFInferenceEngine classes
Update psf_test_setup fixture for reusability and add unit tests for …
43f894f 
…cache handling
Remove unneeded mock of logger in evaluate_model
a88c455
Add wf_psf.inference to list of main packages
6f0e21d
Update version to 3.1.0
cd715c3
docs: Correct syntax error and code-block formatting in doc string ex…
5d54518 
…amples
Remove weights_path references missed during rebase
b5da83f 
Rebase dropped commits that removed weights_path from the metrics interface.
Cleaning up remaining references in docstrings and tests.
Add changelog fragment
98b66e9
@jeipollack
Fix logger formatting for relative RMSE metrics
33ad4f5 
Use f-strings instead of %-formatting to properly display percent symbols
in metric output.
@jeipollack
Update changelog with entry under Bug Fixes
63e6acc
Correct rebase error in configuration.md
0c74bfe
@jeipollack jeipollack force-pushed the feature/159-psf-output-from-trained-model branch from 8ea2ede to 0c74bfe Compare February 2, 2026 15:53
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Reviewers

@ChaitanyaChawak ChaitanyaChawak Awaiting requested review from ChaitanyaChawak

Assignees

@jeipollack jeipollack

Labels

documentation Improvements or additions to documentation enhancement New feature or request

Projects

wavefront-based PSF model estimation
Status: No status

Milestone

EDEN compliance: wf-psf modernization...

Development

Successfully merging this pull request may close these issues.

Add function to output PSF for shape measurement given a trained model

2 participants

@jeipollack @tobias-liaudat