Skip to content

Refactor PSF model optimizer handling with configurable selection and optional TensorFlow Addons#180

Merged
ChaitanyaChawak merged 14 commits intodevelopfrom
feature/179/configurable-optimizer
Feb 2, 2026
Merged

Refactor PSF model optimizer handling with configurable selection and optional TensorFlow Addons#180
ChaitanyaChawak merged 14 commits intodevelopfrom
feature/179/configurable-optimizer

Conversation

@jeipollack
Copy link
Contributor

@jeipollack jeipollack commented Nov 12, 2025

This PR addresses the issue of TensorFlow Addons functions (e.g. RectifiedAdam and interpolate_spline) being incompatible with the EDEN 3.1 (TF 2.11) environment by removing TFA as a required dependency and introducing a fully configurable optimizer selection system.

Closes #179

Key changes:

  • New optimizer.py module with get_optimizer function to centralise optimizer creation.
  • get_optimizer supports multiple input types:
    • RecursiveNamespace from configs,
    • dict of hyperparameters, or
    • string optimizer names ("Adam" or "rectified_adam").
  • Optional RectifiedAdam: dynamically imports TensorFlow Addons only if rectified_adam is defined in the config file.
  • Legacy TensorFlow handling: automatically uses tf.keras.optimizers.legacy.Adam for TF < 2.11 and standard Adam otherwise.
  • Supports hyperparameter overrides: learning rate, beta1/beta2, epsilon, amsgrad.
  • Updated build_PSF_model to accept either a Keras optimizer instance or a config passed through get_optimizer.
  • Removed TFA from required dependencies; training now runs on TF 2.11 without Addons.
  • Optional optimizer parameters configurable via YAML or programmatically for reproducibility.
  • New interpolation.py and types.py modules with copied code from the TFA repo
  • Replaced tfa.image.interpolate_spline with local equivalent renamed totfa_interpolate_spline_rbf
  • Added unit tests:
    • test_optimizer.py: covers legacy/non-legacy handling, RectifiedAdam import logic, overrides, and input normalization.
    • test_interpolation.py: basic validation for the local tfa_interpolate_spline_rbf wrapper.
  • Updates to README, License and new THIRD_PARTY_LICENSE directory with the TFA license
  • Fully backwards compatible; existing unit tests remain valid

@jeipollack jeipollack self-assigned this Nov 12, 2025
@jeipollack jeipollack added the enhancement New feature or request label Nov 12, 2025
@jeipollack jeipollack linked an issue Nov 12, 2025 that may be closed by this pull request
9 tasks
@jeipollack jeipollack force-pushed the feature/179/configurable-optimizer branch from 4cffe99 to 1a53193 Compare November 13, 2025 16:23
@jeipollack jeipollack added the documentation Improvements or additions to documentation label Nov 20, 2025
@jeipollack jeipollack force-pushed the feature/179/configurable-optimizer branch from 7310e49 to ab7d05e Compare January 23, 2026 15:31
@jeipollack jeipollack changed the title Draft: Refactor PSF model optimizer handling with configurable selection and optional TensorFlow Addons Refactor PSF model optimizer handling with configurable selection and optional TensorFlow Addons Jan 23, 2026
@ChaitanyaChawak ChaitanyaChawak merged commit 8360699 into develop Feb 2, 2026
2 checks passed
@ChaitanyaChawak ChaitanyaChawak deleted the feature/179/configurable-optimizer branch February 2, 2026 14:48
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

documentation Improvements or additions to documentation enhancement New feature or request

Projects

Development

Successfully merging this pull request may close these issues.

Remove TensorFlow Addons dependency and introduce configurable optimizer

2 participants