codes.finetune package

codes.finetune package#

Submodules#

codes.finetune.finetune_fcts module#

codes.finetune.finetune_fcts.eval_step_sizes(dt_range, n, rng=None)#

Step sizes to score a model on, one per log spaced bin of dt_range.

Without rng the bin centres are returned, so scores stay comparable. With rng one step size is drawn uniformly in log dt within each bin, which keeps a search from fitting a fixed set of step sizes while keeping coverage even.

Parameters:

  • dt_range – Smallest and largest step size.

  • n (int) – Number of step sizes.

  • rng (Generator | None) – Source of the draws.

Returns:

n step sizes.

Return type:

np.ndarray

codes.finetune.finetune_fcts.finetune_config(conf)#

Resolve the finetune block of a config against its defaults and check it.

Parameters:

conf (dict) – The full CODES config containing a “finetune” block.

Returns:

The finetune settings with every key present.

Return type:

dict

codes.finetune.finetune_fcts.finetune_model(model, train, train_params, timesteps, ft, rng)#

Finetune a trained surrogate for autoregressive rollout.

Each batch combines the ordinary one shot loss with a rollout loss. For the rollout, the model first steps without gradients from the true initial state to a random elapsed time, so the graded window starts on a state the model reached itself. It then takes rollout_steps graded steps, each with its own step size drawn log uniformly from dt_range, and every step is scored against the reference interpolated at the time reached.

Parameters:

  • model – Loaded surrogate implementing forward_from_state.

  • train (ndarray) – Training trajectories, normalised.

  • train_params (ndarray | None) – Training parameters.

  • timesteps (ndarray) – Times the data is stored on.

  • ft (dict) – Finetune settings, see finetune_config.

  • rng (Generator) – Source of all random draws.

Returns:

Per epoch means of the data, rollout and total loss, and the number of batches skipped because their loss was not finite.

Return type:

list[dict]

codes.finetune.finetune_fcts.load_finetune_data(conf)#

Load the train and test splits the way training does.

Finetuning trains on the train split and scores on the test split, which leaves the validation split to run_eval.

Parameters:

conf (dict) – The full CODES config.

Returns:

Arrays “train”, “test”, “train_params”, “test_params”, “timesteps”.

Return type:

dict

codes.finetune.finetune_fcts.load_trained_model(surr_name, conf, data, device)#

Build a surrogate and load the main model of the configured training run.

Parameters:

  • surr_name (str) – Name of the surrogate class.

  • conf (dict) – The full CODES config.

  • data (dict) – Output of load_finetune_data.

  • device (str) – Device to place the model on.

Returns:

The loaded surrogate, in eval mode.

codes.finetune.finetune_fcts.rollout_errors(model, data, params, timesteps, dt, batch_size=1024)#

Absolute errors of a fixed step rollout to the stored horizon.

Matches the fixed_step mode of the iterative evaluation. Errors are taken in the units of the *_log metrics: log10 when log10_transform is on, otherwise the units of the data.

Parameters:

  • model – Loaded surrogate implementing forward_from_state.

  • data (ndarray) – Trajectories, normalised.

  • params (ndarray | None) – Parameters.

  • timesteps (ndarray) – Times the data is stored on.

  • dt (float) – Step size, in the units of timesteps.

  • batch_size (int) – Trajectories per forward pass.

Returns:

Errors of shape (n_samples, n_steps, n_quantities).

Return type:

ndarray

codes.finetune.finetune_fcts.rollout_score(model, data, params, timesteps, step_sizes, dynamic_span=None, batch_size=1024)#

Score a model by its rollout error over several step sizes.

For each step size the median absolute error is taken over all trajectories, rollout steps and quantities. The score is the geometric mean of these medians, which weights every step size equally even when their errors differ by orders of magnitude. Non finite errors count as NONFINITE_ERROR.

Parameters:

  • model – Loaded surrogate implementing forward_from_state.

  • data (ndarray) – Trajectories, normalised.

  • params (ndarray | None) – Parameters.

  • timesteps (ndarray) – Times the data is stored on.

  • step_sizes – Step sizes, in the units of timesteps.

  • dynamic_span (float | None) – If set, score only trajectories in which some quantity changes by at least this much over the horizon, in the units of the errors.

  • batch_size (int) – Trajectories per forward pass.

Returns:

The score and the median error per step size.

Return type:

tuple[float, list[float]]

codes.finetune.finetune_fcts.run_finetuning(conf, device)#

Finetune the main model of every configured surrogate for rollout.

The finetuned models are saved as an ordinary training run under trained/<finetune.training_id>, with a config that run_eval can use. A summary of the scores before and after is written to results/<finetune.training_id>.

Parameters:

  • conf (dict) – The full CODES config with a finetune block.

  • device (str) – Device to train on.

Returns:

Per surrogate summary of the finetuning.

Return type:

dict

Module contents#

codes.finetune.eval_step_sizes(dt_range, n, rng=None)#

Step sizes to score a model on, one per log spaced bin of dt_range.

Without rng the bin centres are returned, so scores stay comparable. With rng one step size is drawn uniformly in log dt within each bin, which keeps a search from fitting a fixed set of step sizes while keeping coverage even.

Parameters:

  • dt_range – Smallest and largest step size.

  • n (int) – Number of step sizes.

  • rng (Generator | None) – Source of the draws.

Returns:

n step sizes.

Return type:

np.ndarray

codes.finetune.finetune_config(conf)#

Resolve the finetune block of a config against its defaults and check it.

Parameters:

conf (dict) – The full CODES config containing a “finetune” block.

Returns:

The finetune settings with every key present.

Return type:

dict

codes.finetune.finetune_model(model, train, train_params, timesteps, ft, rng)#

Finetune a trained surrogate for autoregressive rollout.

Each batch combines the ordinary one shot loss with a rollout loss. For the rollout, the model first steps without gradients from the true initial state to a random elapsed time, so the graded window starts on a state the model reached itself. It then takes rollout_steps graded steps, each with its own step size drawn log uniformly from dt_range, and every step is scored against the reference interpolated at the time reached.

Parameters:

  • model – Loaded surrogate implementing forward_from_state.

  • train (ndarray) – Training trajectories, normalised.

  • train_params (ndarray | None) – Training parameters.

  • timesteps (ndarray) – Times the data is stored on.

  • ft (dict) – Finetune settings, see finetune_config.

  • rng (Generator) – Source of all random draws.

Returns:

Per epoch means of the data, rollout and total loss, and the number of batches skipped because their loss was not finite.

Return type:

list[dict]

codes.finetune.load_finetune_data(conf)#

Load the train and test splits the way training does.

Finetuning trains on the train split and scores on the test split, which leaves the validation split to run_eval.

Parameters:

conf (dict) – The full CODES config.

Returns:

Arrays “train”, “test”, “train_params”, “test_params”, “timesteps”.

Return type:

dict

codes.finetune.load_trained_model(surr_name, conf, data, device)#

Build a surrogate and load the main model of the configured training run.

Parameters:

  • surr_name (str) – Name of the surrogate class.

  • conf (dict) – The full CODES config.

  • data (dict) – Output of load_finetune_data.

  • device (str) – Device to place the model on.

Returns:

The loaded surrogate, in eval mode.

codes.finetune.rollout_errors(model, data, params, timesteps, dt, batch_size=1024)#

Absolute errors of a fixed step rollout to the stored horizon.

Matches the fixed_step mode of the iterative evaluation. Errors are taken in the units of the *_log metrics: log10 when log10_transform is on, otherwise the units of the data.

Parameters:

  • model – Loaded surrogate implementing forward_from_state.

  • data (ndarray) – Trajectories, normalised.

  • params (ndarray | None) – Parameters.

  • timesteps (ndarray) – Times the data is stored on.

  • dt (float) – Step size, in the units of timesteps.

  • batch_size (int) – Trajectories per forward pass.

Returns:

Errors of shape (n_samples, n_steps, n_quantities).

Return type:

ndarray

codes.finetune.rollout_score(model, data, params, timesteps, step_sizes, dynamic_span=None, batch_size=1024)#

Score a model by its rollout error over several step sizes.

For each step size the median absolute error is taken over all trajectories, rollout steps and quantities. The score is the geometric mean of these medians, which weights every step size equally even when their errors differ by orders of magnitude. Non finite errors count as NONFINITE_ERROR.

Parameters:

  • model – Loaded surrogate implementing forward_from_state.

  • data (ndarray) – Trajectories, normalised.

  • params (ndarray | None) – Parameters.

  • timesteps (ndarray) – Times the data is stored on.

  • step_sizes – Step sizes, in the units of timesteps.

  • dynamic_span (float | None) – If set, score only trajectories in which some quantity changes by at least this much over the horizon, in the units of the errors.

  • batch_size (int) – Trajectories per forward pass.

Returns:

The score and the median error per step size.

Return type:

tuple[float, list[float]]

codes.finetune.run_finetuning(conf, device)#

Finetune the main model of every configured surrogate for rollout.

The finetuned models are saved as an ordinary training run under trained/<finetune.training_id>, with a config that run_eval can use. A summary of the scores before and after is written to results/<finetune.training_id>.

Parameters:

  • conf (dict) – The full CODES config with a finetune block.

  • device (str) – Device to train on.

Returns:

Per surrogate summary of the finetuning.

Return type:

dict