Coverage for pySDC/tutorial/step_7/torch_heat.py: 79%
133 statements
« prev ^ index » next coverage.py v7.17.0, created at 2026-11-02 12:30 +0000
1"""
2PyTorch helpers for step 7 part D: a tensor datatype, a heat equation problem using it and a (too) simple network.
3"""
5import numpy as np
6import scipy.sparse as sp
7import torch
8import torch.nn as nn
9from pySDC.core.errors import ProblemError
10from pySDC.core.problem import Problem, WorkCounter
11from pySDC.helpers import problem_helper
13try:
14 from mpi4py import MPI
15except ImportError:
16 MPI = None
19class Tensor(torch.Tensor):
20 """
21 Wrapper for PyTorch tensor.
22 Be aware that this is totally WIP! Should be fine to count iterations, but desperately needs cleaning up if this
23 project goes much further!
25 TODO: Have to update `torch/multiprocessing/reductions.py` in order to share this datatype across processes.
27 Attributes:
28 comm: MPI communicator or None
29 """
31 comm = None
33 @staticmethod
34 def __new__(cls, init, val=0.0, *args, **kwargs):
35 """
36 Instantiates new datatype. This ensures that even when manipulating data, the result is still a tensor.
38 Args:
39 init: either another mesh or a tuple containing the dimensions, the communicator and the dtype
40 val: value to initialize
42 Returns:
43 obj of type mesh
45 """
46 # TODO: The cloning of tensors going in is likely slow
48 if isinstance(init, torch.Tensor):
49 obj = super().__new__(cls, init.clone())
50 obj[:] = init[:]
51 elif (
52 isinstance(init, tuple)
53 and (init[1] is None or isinstance(init[1], MPI.Intracomm))
54 # and isinstance(init[2], np.dtype)
55 ):
56 if isinstance(init[0][0], torch.Tensor):
57 obj = super().__new__(cls, init[0].clone())
58 else:
59 obj = super().__new__(cls, *init[0])
60 obj.fill_(val)
61 cls.comm = init[1]
62 else:
63 raise NotImplementedError(type(init))
64 return obj
66 def __abs__(self):
67 """
68 Overloading the abs operator
70 Returns:
71 float: absolute maximum of all mesh values
72 """
73 # take absolute values of the mesh values
74 local_absval = float(torch.amax(torch.abs(self)))
76 if self.comm is not None:
77 if self.comm.Get_size() > 1:
78 global_absval = 0.0
79 global_absval = max(self.comm.allreduce(sendobj=local_absval, op=MPI.MAX), global_absval)
80 else:
81 global_absval = local_absval
82 else:
83 global_absval = local_absval
85 return float(global_absval)
87 def isend(self, dest=None, tag=None, comm=None):
88 """
89 Routine for sending data forward in time (non-blocking)
91 Args:
92 dest (int): target rank
93 tag (int): communication tag
94 comm: communicator
96 Returns:
97 request handle
98 """
99 return comm.Issend(self[:], dest=dest, tag=tag)
101 def irecv(self, source=None, tag=None, comm=None):
102 """
103 Routine for receiving in time
105 Args:
106 source (int): source rank
107 tag (int): communication tag
108 comm: communicator
110 Returns:
111 None
112 """
113 return comm.Irecv(self[:], source=source, tag=tag)
115 def bcast(self, root=None, comm=None):
116 """
117 Routine for broadcasting values
119 Args:
120 root (int): process with value to broadcast
121 comm: communicator
123 Returns:
124 broadcasted values
125 """
126 comm.Bcast(self[:], root=root)
127 return self
130class HeatEquationModel(nn.Module):
131 """
132 Very simple model to learn the heat equation. Beware! It's too simple.
133 Some machine learning expert please fix this!
134 """
136 def __init__(self, problem, hidden_size=64):
137 self.input_size = problem.nvars * 3
138 self.output_size = problem.nvars
139 self.problem = problem
141 super().__init__()
143 self.fc1 = nn.Linear(self.input_size, hidden_size)
144 self.relu = nn.ReLU()
145 self.fc2 = nn.Linear(hidden_size, self.output_size)
147 # Initialize weights (example)
148 nn.init.xavier_uniform_(self.fc1.weight)
149 nn.init.xavier_uniform_(self.fc2.weight)
151 def forward(self, x, t, dt):
152 # prepare individual tensors
153 x = x.float()
154 _t = torch.ones(x.shape) * dt
155 _dt = torch.ones(x.shape) * dt
157 # Concatenate t and dt with the input x
158 _x = torch.cat((x, _t, _dt), dim=0)
160 _x = self.fc1(_x)
161 _x = self.relu(_x)
162 _x = self.fc2(_x)
163 return _x
165 def __call__(self, *args, **kwargs):
166 me = self.problem.u_init
167 me[:] = super().__call__(*args, **kwargs)
168 return me
171class Heat1DFDTensor(Problem):
172 """
173 Very simple 1-dimensional finite differences implementation of a heat equation using the pySDC-PyTorch interface.
174 Still includes some mess.
175 """
177 dtype_u = Tensor
178 dtype_f = Tensor
180 def __init__(
181 self,
182 nvars=256,
183 nu=1.0,
184 freq=4,
185 stencil_type='center',
186 order=2,
187 lintol=1e-12,
188 liniter=10000,
189 solver_type='direct',
190 bc='periodic',
191 bcParams=None,
192 ):
193 # make sure parameters have the correct types
194 if type(nvars) not in [int, tuple]:
195 raise ProblemError('nvars should be either tuple or int')
196 if type(freq) not in [int, tuple]:
197 raise ProblemError('freq should be either tuple or int')
199 ndim = 1
201 # eventually extend freq to other dimension
202 if type(freq) is int:
203 freq = (freq,) * ndim
204 if len(freq) != ndim:
205 raise ProblemError(f'len(freq)={len(freq)}, different to ndim={ndim}')
207 # check values for freq and nvars
208 for f in freq:
209 if ndim == 1 and f == -1:
210 # use Gaussian initial solution in 1D
211 bc = 'periodic'
212 break
213 if f % 2 != 0 and bc == 'periodic':
214 raise ProblemError('need even number of frequencies due to periodic BCs')
216 # invoke super init, passing number of dofs
217 super().__init__(init=(torch.empty(size=(nvars,), dtype=torch.double), None, np.dtype('float64')))
219 dx, xvalues = problem_helper.get_1d_grid(size=nvars, bc=bc, left_boundary=0.0, right_boundary=1.0)
221 self.A_, _ = problem_helper.get_finite_difference_matrix(
222 derivative=2,
223 order=order,
224 stencil_type=stencil_type,
225 dx=dx,
226 size=nvars,
227 dim=ndim,
228 bc=bc,
229 )
230 self.A_ *= nu
231 self.A = torch.tensor(self.A_.todense())
233 self.xvalues = torch.tensor(xvalues, dtype=torch.double)
234 self.Id = torch.tensor((sp.eye(nvars, format='csc')).todense())
236 # store attribute and register them as parameters
237 self._makeAttributeAndRegister('nvars', 'stencil_type', 'order', 'bc', 'nu', localVars=locals(), readOnly=True)
238 self._makeAttributeAndRegister('freq', 'lintol', 'liniter', 'solver_type', localVars=locals())
240 if self.solver_type != 'direct':
241 self.work_counters[self.solver_type] = WorkCounter()
243 @property
244 def ndim(self):
245 """Number of dimensions of the spatial problem"""
246 return 1
248 @property
249 def dx(self):
250 """Size of the mesh (in all dimensions)"""
251 return self.xvalues[1] - self.xvalues[0]
253 @property
254 def grids(self):
255 """ND grids associated to the problem"""
256 x = self.xvalues
257 if self.ndim == 1:
258 return x
259 if self.ndim == 2:
260 return x[None, :], x[:, None]
261 if self.ndim == 3:
262 return x[None, :, None], x[:, None, None], x[None, None, :]
264 def eval_f(self, u, t):
265 """
266 Routine to evaluate the right-hand side of the problem.
268 Parameters
269 ----------
270 u : dtype_u
271 Current values.
272 t : float
273 Current time.
275 Returns
276 -------
277 f : dtype_f
278 Values of the right-hand side of the problem.
279 """
280 f = self.f_init
281 f[:] = torch.matmul(self.A, u)
282 return f
284 def ML_predict(self, u0, t0, dt):
285 """
286 Predict the solution at t0+dt given initial conditions u0
287 """
288 # read in model
289 model = HeatEquationModel(self)
290 model.load_state_dict(torch.load('heat_equation_model.pth'))
291 model.eval()
293 # evaluate model
294 predicted_state = model(u0, t0, dt)
295 sol = self.u_init
296 sol[:] = predicted_state.double()[:]
297 return sol
299 def solve_system(self, rhs, factor, u0, t):
300 r"""
301 Simple linear solver for :math:`(I-factor\cdot A)\vec{u}=\vec{rhs}`.
303 Parameters
304 ----------
305 rhs : dtype_f
306 Right-hand side for the linear system.
307 factor : float
308 Abbrev. for the local stepsize (or any other factor required).
309 u0 : dtype_u
310 Initial guess for the iterative solver.
311 t : float
312 Current time (e.g. for time-dependent BCs).
314 Returns
315 -------
316 sol : dtype_u
317 The solution of the linear solver.
318 """
319 solver_type, Id, A, nvars, sol = (
320 self.solver_type,
321 self.Id,
322 self.A,
323 self.nvars,
324 self.u_init,
325 )
327 if solver_type == 'direct':
328 sol[:] = torch.linalg.solve(Id - factor * A, rhs.flatten()).reshape(nvars)
329 # TODO: implement torch equivalent of cg
330 # elif solver_type == 'CG':
331 # sol[:] = cg(
332 # Id - factor * A,
333 # rhs.flatten(),
334 # x0=u0.flatten(),
335 # tol=lintol,
336 # maxiter=liniter,
337 # atol=0,
338 # callback=self.work_counters[solver_type],
339 # )[0].reshape(nvars)
340 else:
341 raise ValueError(f'solver type "{solver_type}" not known!')
343 return sol
345 def u_exact(self, t, **kwargs):
346 r"""
347 Routine to compute the exact solution at time :math:`t`.
349 Parameters
350 ----------
351 t : float
352 Time of the exact solution.
354 Returns
355 -------
356 sol : dtype_u
357 The exact solution.
358 """
359 if 'u_init' in kwargs.keys() or 't_init' in kwargs.keys():
360 self.logger.warning(
361 f'{type(self).__name__} uses an analytic exact solution from t=0. If you try to compute the local error, you will get the global error instead!'
362 )
364 ndim, freq, nu, dx, sol = self.ndim, self.freq, self.nu, self.dx, self.u_init
366 if ndim == 1:
367 x = self.grids
368 rho = (2.0 - 2.0 * torch.cos(np.pi * freq[0] * dx)) / dx**2
369 if freq[0] > 0:
370 sol[:] = torch.sin(np.pi * freq[0] * x) * torch.exp(-t * nu * rho)
371 else:
372 raise NotImplementedError
374 return sol