Fix division by zero in math operations

[Ayushhgit]

Describe your changes

Summary of the changes:

• Added a 1e-6 epsilon to the output of torch.nn.functional.softplus in [base_graph_model.py].
• Added a 1e-8 epsilon to the self.feature_weights inverse square root calculation inside ARModel.__init__.

Motivation and context:
These changes resolve two separate division-by-zero crashes:

1. When computing the pred_std using softplus, inputs $\le -80$ in float32 will evaluate exactly to 0.0. This zero denominator causes cascading NaN outputs instantly destroying training across metrics like wmse or wmae.
2. When users specify a manual feature weight of 0.0 for specific state variables they want the model to ignore, creating per_var_std triggers a scalar division by zero, crashing model initialization. The epsilons bring numerical stability against these mathematical singularities.

Dependencies:
None

Issue Link

N/A

Type of change

• 🐛 Bug fix (non-breaking change that fixes an issue)
• ✨ New feature (non-breaking change that adds functionality)
• 💥 Breaking change (fix or feature that would cause existing functionality to not work as expected)
• 📖 Documentation (Addition or improvements to documentation)

Checklist before requesting a review

• My branch is up-to-date with the target branch - if not update your fork with the changes from the target branch (use pull with --rebase option if possible).
• I have performed a self-review of my code
• For any new/modified functions/classes I have added docstrings that clearly describe its purpose, expected inputs and returned values
• I have placed in-line comments to clarify the intent of any hard-to-understand passages of my code
• I have updated the README to cover introduced code changes
• I have added tests that prove my fix is effective or that my feature works
• I have given the PR a name that clearly describes the change, written in imperative form (context).
• I have requested a reviewer and an assignee (assignee is responsible for merging). This applies only if you have write access to the repo, otherwise feel free to tag a maintainer to add a reviewer and assignee.

Checklist for reviewers

Each PR comes with its own improvements and flaws. The reviewer should check the following:

• the code is readable
• the code is well tested
• the code is documented (including return types and parameters)
• the code is easy to maintain

Author checklist after completed review

• I have added a line to the CHANGELOG describing this change, in a section
  reflecting type of change (add section where missing):
  • added: when you have added new functionality
  • changed: when default behaviour of the code has been changed
  • fixes: when your contribution fixes a bug
  • maintenance: when your contribution is relates to repo maintenance, e.g. CI/CD or documentation

Checklist for assignee

• PR is up to date with the base branch
• the tests pass
• (if the PR is not just maintenance/bugfix) the PR is assigned to the next milestone. If it is not, propose it for a future milestone.
• author has added an entry to the changelog (and designated the change as added, changed, fixed or maintenance)
• Once the PR is ready to be merged, squash commits and merge the PR.

[joeloskarsson]

Thanks, these are good changes to make the code more robust. Would it be possible to also add some test that checks that we don't run into numerical issues with the chosen constants?

[Ayushhgit]

@joeloskarsson Thanks for the feedback! I've just pushed a commit with the updates:

Warning: Added a UserWarning in ar_model.py that alerts the user and lists the indices if any variables receive a feature weight of exactly 0.0 (excluding them from the loss). Tests: Added

test_numerical_stability.py with 5 pure tensor unit tests. These verify that the softplus epsilon (1e-6) and the feature weight epsilon (1e-8) successfully prevent exact zeros and subnormals, keeping downstream loss metrics (wmse, wmae, per_var_std) perfectly finite even with worst-case or zero-weight inputs. Let me know if you need anything else!

[sadamov]

please link this issue in your PR description: #526

[sadamov]

Thanks @Ayushhgit, rebased onto current main and scoped down. Two changes:

1. Dropped the softplus(pred_std_raw) + 1e-6 piece - that line overlaps with #523's * self.diff_std scaling, and once pred_std starts on the empirical scale via #523 the underflow concern is largely moot. Better to land the two ideas separately so they don't fight.
2. Switched to torch.finfo(torch.float32).eps for the sqrt eps (matches the existing state_std / forcing_std epsilon choice a few lines below).

tests requested by @joeloskarsson have been added before. please review

Refs #526. Merging after CI is green.

[joeloskarsson]

Thanks for adding the tests, but I am a bit confused here since they don't call any of the actual code. E.g. mirroring the per_var_std calculations separately in the test file seems very bad to maintain. Testing the numerical aspect is good, but we should do it running through the actual code used when training models. If this code is not structured such that it can be tested (requiring a mirroring function), then we should restructure it so that is possible.