pinnx 0.0.1__py2.py3-none-any.whl

pinnx/__init__.py

@@ -0,0 +1,13 @@
+__version__ = "0.0.1"
+
+from . import callbacks
+from . import data
+from . import geometry
+from . import grad
+from . import icbc
+from . import metrics
+from . import nn
+from . import utils
+from ._model import Model
+from pinnx._convert import array_to_dict, dict_to_array
+from .utils import saveplot

pinnx/_convert.py

@@ -0,0 +1,49 @@
+# Copyright 2024 BDP Ecosystem Limited. 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.
+# ==============================================================================
+
+import brainunit as u
+import numpy as np
+
+__all__ = [
+    "array_to_dict",
+    "dict_to_array",
+]
+
+
+def array_to_dict(x, *names):
+    """
+    Convert args to a dictionary.
+    """
+    if x.shape[-1] != len(names):
+        raise ValueError("The number of columns of x must be equal to the number of names.")
+
+    return {key: x[..., i] for i, key in enumerate(names)}
+
+
+def dict_to_array(d):
+    """
+    Convert a dictionary to an array.
+
+    Args:
+        d (dict): The dictionary.
+
+    Returns:
+        ndarray: The array.
+    """
+    keys = tuple(d.keys())
+    if isinstance(d[keys[0]], np.ndarray):
+        return np.stack([d[key] for key in keys], axis=-1)
+    else:
+        return u.math.stack([d[key] for key in keys], axis=-1)

pinnx/_display.py

@@ -0,0 +1,72 @@
+import sys
+
+import numpy as np
+
+from pinnx.utils import list_to_str
+
+
+class TrainingDisplay:
+    """
+    Display training progress.
+    """
+
+    def __init__(self):
+        self.len_train = None
+        self.len_test = None
+        self.len_metric = None
+        self.is_header_print = False
+
+    def print_one(self, s1, s2, s3, s4):
+        print(
+            "{:{l1}s}{:{l2}s}{:{l3}s}{:{l4}s}".format(
+                s1,
+                s2,
+                s3,
+                s4,
+                l1=10,
+                l2=self.len_train,
+                l3=self.len_test,
+                l4=self.len_metric,
+            )
+        )
+        sys.stdout.flush()
+
+    def header(self):
+        self.print_one("Step", "Train loss", "Test loss", "Test metric")
+        self.is_header_print = True
+
+    def __call__(self, train_state):
+        if not self.is_header_print:
+            self.len_train = len(train_state.loss_train) * 10 + 4
+            self.len_test = len(train_state.loss_test) * 10 + 4
+            self.len_metric = len(train_state.metrics_test) * 10 + 4
+            self.header()
+        self.print_one(
+            str(train_state.step),
+            list_to_str(train_state.loss_train),
+            list_to_str(train_state.loss_test),
+            list_to_str(train_state.metrics_test),
+        )
+
+    def summary(self, train_state):
+        print("Best model at step {:d}:".format(train_state.best_step))
+        print("  train loss: {:.2e}".format(train_state.best_loss_train))
+        print("  test loss: {:.2e}".format(train_state.best_loss_test))
+        print("  test metric: {:s}".format(list_to_str(train_state.best_metrics)))
+        if train_state.best_ystd is not None:
+            print("  Uncertainty:")
+            print("    l2: {:g}".format(np.linalg.norm(train_state.best_ystd)))
+            print(
+                "    l_infinity: {:g}".format(
+                    np.linalg.norm(train_state.best_ystd, ord=np.inf)
+                )
+            )
+            print(
+                "    max uncertainty location:",
+                train_state.X_test[np.argmax(train_state.best_ystd)],
+            )
+        print("")
+        self.is_header_print = False
+
+
+training_display = TrainingDisplay()