🚀[FEA]: Time (and Space) Marching Inferencer

[androbomb]

Is this a new feature, an improvement, or a change to existing functionality?

New Feature

How would you describe the priority of this feature request

Low (would be nice)

Please provide a clear description of problem you would like to solve.

I feel user should have the possibility to add both Time and Space marching schedule (similarly to what done in the Wave1D example, but unfortunately no code was given, 1D Wave Equation - NVIDIA Docs ).

The idea is to have the possibility to add a custom training step-dependent parametrization/criteria of a region, so that it updates during training, e.g., we could pass to the PointwiseInteriorConstraint a lambda function which returns a certain sympy constraint/parametrisation.

This could help the training when involving cases with transient solutions and localised sources.

Describe any alternatives you have considered

I have implemented a custom PointwiseConstraint, extending the Constraint class, by updating criteria/parametrization every epoch before computing losses, e.g.

def loss(self, step: int) -> Dict[str, torch.Tensor]:
    """
    Loss function. It is here that the whole change enters.
    
    We need to update the dataset before calling _loss method.
    """
    if self._output_vars is None:
        logger.warn("Calling loss without forward call")
        return {}
    
    # call update dataset
    self.update_dataset(step = step) ## <==== HERE ================
    
    losses = self._loss(
        self._input_vars,
        self._output_vars,
        self._target_vars,
        self._lambda_weighting,
        step,
    )
    
    return losses

Where self.update_dataset(step) updates the dataset and the dataloader by resampling the geometry with the updated parametrization/criteria; e.g. for the time-marching schedule of wave example I could pass something like

time_marching_max_steps = min(cfg.training.max_steps, cfg.custom.time_marching_max_steps)
_time_marching_schedule = lambda step: {
    t_symbol : ( 
        _t_i, 
        min(_t_i + step / time_marching_max_steps  * _t_f, _t_f) 
    )
}

and similarly I could do something with the space marching, e.g, for the Wave1D example mentioned above,

_max_marching_steps = 60*cfg.training.max_steps // 100 # 80% of the training epochs
_base_spatial_size = L/8
_spatial_marching_schedule = lambda step: And(
    (Symbol("x") >= max(0, L/2 - ( _base_spatial_size + step/_max_marching_steps) )  ),
    (Symbol("x") <= min(L, L/2 + ( _base_spatial_size + step/_max_marching_steps) )  )
)

And passing it to the class as

interior = PointwiseInteriorExpandingConstraint(
    nodes    = nodes,
    geometry = geo,
    outvar   = {"wave_equation": 0},
    batch_size = cfg.batch_size.interior,
    parameterization  = time_range,
    lambda_weighting={"wave_equation": 100.0}, 
    # Adds on
    expanding_criteria               = _spatial_marching_schedule,
    expanding_parametrization = _time_marching_schedule,
)
domain.add_constraint(interior, "interior")

Additional context

No response