feat(callback): implement a MonitorCallback class
using MonitorCallback is more flexible than passing a monitor and supports saving visualizations to disk
+import os
import sys
+import dill
+import warnings
+import logging
import torch
import torch.optim as optim
import torch.nn as nn
@@ -15,7 +19,7 @@ from .version_utils import warn_deprecate_class
from .generator import Generator3D, GeneratorSpherical
from inspect import signature
from copy import deepcopy
-import warnings
+from datetime import datetime
# the name ExampleGenerator3D is deprecated
ExampleGenerator3D = warn_deprecate_class(Generator3D)
@@ -1177,3 +1181,45 @@ class MonitorSphericalHarmonics(MonitorSpherical):
warnings.warn(f"Error caught when accessing {self.__class__.__name__}, returning None:\n{e}")
ret = None
return ret
+
+
+# callbacks to be passed to SphericalSolver.fit()
+
+
+class MonitorCallback:
+ def __init__(self, monitor, fig_dir=None, check_against='local', repaint_last=True):
+ self.monitor = monitor
+ self.fig_dir = fig_dir
+ self.repaint_last = repaint_last
+ if check_against not in ['local', 'global']:
+ raise ValueError(f'unknown check_against type = {check_against}')
+ self.check_against = check_against
+
+ def to_repaint(self, solver):
+ if self.check_against == 'local':
+ epoch_now = solver.local_epoch + 1
+ elif self.check_against == 'global':
+ epoch_now = solver.global_epoch + 1
+ else:
+ raise ValueError(f'unknown check_against type = {self.check_against}')
+
+ if epoch_now % self.monitor.check_every == 0:
+ return True
+ if self.repaint_last and solver.local_epoch == solver._max_local_epoch - 1:
+ return True
+
+ return False
+
+ def __call__(self, solver):
+ if not self.to_repaint(solver):
+ return
+
+ self.monitor.check(
+ solver.nets,
+ solver.conditions,
+ loss_history=solver.loss,
+ analytic_mse_history=solver.analytic_solutions
+ )
+ if self.fig_dir:
+ pic_path = os.path.join(self.fig_dir, f"epoch-{solver.global_epoch}.png")
+ self.monitor.fig.savefig(pic_path)
\ No newline at end of file
