Source code for modulus.sym.models.moving_time_window

# Copyright (c) 2023, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# See the License for the specific language governing permissions and
# limitations under the License.

from typing import Dict
from modulus.sym.key import Key
import copy

import torch
import torch.nn as nn
from torch import Tensor

from modulus.sym.models.arch import Arch

[docs]class MovingTimeWindowArch(Arch): """ Moving time window model the keeps track of current time window and previous window. Parameters ---------- arch : Arch Modulus architecture to use for moving time window. window_size : float Size of the time window. This will be used to slide the window forward every iteration. """ def __init__( self, arch: Arch, window_size: float, ) -> None: output_keys = ( arch.output_keys + [Key( + "_prev_step") for x in arch.output_keys] + [Key( + "_prev_step_diff") for x in arch.output_keys] ) super().__init__( input_keys=arch.input_keys, output_keys=output_keys, periodicity=arch.periodicity, ) # set networks for current and prev time window self.arch_prev_step = arch self.arch = copy.deepcopy(arch) # store time window parameters self.window_size = window_size self.window_location = nn.Parameter(torch.empty(1), requires_grad=False) self.reset_parameters()
[docs] def forward(self, in_vars: Dict[str, Tensor]) -> Dict[str, Tensor]: with torch.no_grad(): in_vars["t"] += self.window_location y_prev_step = self.arch_prev_step.forward(in_vars) y = self.arch.forward(in_vars) y_keys = list(y.keys()) for key in y_keys: y_prev = y_prev_step[key] y[key + "_prev_step"] = y_prev y[key + "_prev_step_diff"] = y[key] - y_prev return y

def move_window(self): += self.window_size for param, param_prev_step in zip( self.arch.parameters(), self.arch_prev_step.parameters() ): = param.detach().clone().data param_prev_step.requires_grad = False def reset_parameters(self) -> None: nn.init.constant_(self.window_location, 0)

© Copyright 2023, NVIDIA Modulus Team. Last updated on Jan 25, 2024.