# SPDX-FileCopyrightText: Copyright (c) 2023 - 2024 NVIDIA CORPORATION & AFFILIATES.
# SPDX-FileCopyrightText: All rights reserved.
# SPDX-License-Identifier: Apache-2.0
#
# 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.
"""base class for PDEs"""
from sympy import (
Function,
init_printing,
preview,
Matrix,
Eq,
)
from typing import Dict, List
from physicsnemo.sym.node import Node
[docs]
class PDE(object):
"""base class for all partial differential equations"""
name = "PDE"
def __init__(self):
super().__init__()
self.equations = Variables()
[docs]
def pprint(self, print_latex=False):
"""
Print differential equation.
Parameters
----------
print_latex : bool
If True print the equations in Latex. Else, just
print as text.
"""
init_printing(use_latex=True)
for key, value in self.equations.items():
print(str(key) + ": " + str(value))
if print_latex:
preview(
Matrix(
[
Eq(Function(name, real=True), eq)
for name, eq in self.equations.items()
]
),
mat_str="cases",
mat_delim="",
)
def subs(self, x, y):
for name, eq in self.equations.items():
self.equations[name] = eq.subs(x, y).doit()
[docs]
def make_nodes(
self,
create_instances: int = 1,
freeze_terms: Dict[str, List[int]] = None,
detach_names: list[str] = None,
return_as_dict: bool = False,
):
"""
Make a list of nodes from PDE.
Parameters
----------
create_instances : int
This will create various instances of the same equations
freeze_terms : Dict[str, List[int]]
This will freeze the terms in appropiate equation
detach_names : List[str]
This will detach the inputs of the resulting node.
return_as_dict : bool
If True, return nodes as a dictionary with equation names as keys.
If False, return nodes as a list (default behavior).
Returns
-------
nodes : list[Node] | dict[str, Node]
Makes a separate node for every equation.
Returns list of nodes if return_as_dict=False, dictionary if return_as_dict=True.
"""
if detach_names is None:
detach_names = []
if freeze_terms is None:
freeze_terms = {}
if create_instances == 1:
if bool(freeze_terms):
print(
"Freezing of terms is not supported when create_instance = 1. No terms will be frozen!"
)
freeze_terms = {} # override with an empty dict
node_dict = {
str(name): Node.from_sympy(eq, str(name), freeze_terms, detach_names)
for name, eq in self.equations.items()
}
else:
# look for empty lists in freeze_terms dict
freeze_terms = {k: v for k, v in freeze_terms.items() if v}
node_dict = {}
for i in range(create_instances):
for name, eq in self.equations.items():
node_name = f"{name}_{i}"
if node_name in freeze_terms.keys():
node = Node.from_sympy(
eq,
node_name,
freeze_terms[node_name],
detach_names,
)
else:
# set the freeze terms to an empty list
print(
f"No freeze terms found for instance {node_name}, setting to empty."
)
node = Node.from_sympy(
eq,
node_name,
[],
detach_names,
)
node_dict[node_name] = node
if return_as_dict:
return node_dict
else:
return list(node_dict.values())
