deeplearning/modulus/modulus-sym-v130/_modules/modulus/sym/models/moving_time_window.html
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
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# 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(x.name + "_prev_step") for x in arch.output_keys]
+ [Key(x.name + "_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 ydef move_window(self):
self.window_location.data += self.window_size
for param, param_prev_step in zip(
self.arch.parameters(), self.arch_prev_step.parameters()
):
param_prev_step.data = param.detach().clone().data
param_prev_step.requires_grad = False
def reset_parameters(self) -> None:
nn.init.constant_(self.window_location, 0)