File size: 1,906 Bytes
338c3e4 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 | from torch import nn, Tensor
from typing import Optional, List
from .loss import MseLoss
class CfdModel(nn.Module):
"""
Base model for all data-driven NN for CFD, learns the mapping from
conditions (physics properties, boundary conditions and geometry)
to the solution at a later time.
"""
def __init__(self, loss_fn: MseLoss):
super().__init__()
self.loss_fn = loss_fn
def forward(
self,
x: Tensor,
mask: Optional[Tensor] = None,
label: Optional[Tensor] = None,
case_params: Optional[dict] = None,
) -> dict:
raise NotImplementedError
def generate_one(
self,
case_params,
t: Tensor,
height: int,
width: int,
**kwargs,
) -> Tensor:
"""
Generate the frame at time step `time`, given the case parameters.
`case_params`.
"""
raise NotImplementedError
class AutoCfdModel(nn.Module):
"""
A CFD model that generates the solution auto-regressively, one frame
at a time.
"""
def __init__(self, loss_fn: nn.Module):
super().__init__()
self.loss_fn = loss_fn
def forward(
self,
inputs: Tensor,
label: Optional[Tensor] = None,
case_params: Optional[dict] = None,
mask: Optional[Tensor] = None,
**kwargs,
) -> dict:
raise NotImplementedError
def generate(
self,
inputs: Tensor,
case_params: Tensor,
mask: Tensor,
**kwargs,
) -> Tensor:
raise NotImplementedError
def generate_many(
self,
inputs: Tensor,
case_params: Tensor,
mask: Tensor,
steps: int,
**kwargs,
) -> List[Tensor]:
"""
Given a frame `inputs`, generate the next `steps` frames.
"""
raise NotImplementedError
|