xref: /petsc/src/binding/petsc4py/src/petsc4py/typing.py (revision 3e2ee6483bc6ed4e19bf12bb0af887c748eb13d9) 
1# Author:  Lisandro Dalcin
2# Contact: dalcinl@gmail.com
3"""Typing support."""
4
5from __future__ import annotations  # novermin
6from typing import (  # novermin
7    Callable,
8    Sequence,
9    Literal,
10    TypeAlias,
11)
12from numpy.typing import (
13    NDArray,
14)
15import numpy as np
16from .PETSc import (
17    InsertMode,
18    ScatterMode,
19    NormType,
20    Object,
21    Vec,
22    Mat,
23    NullSpace,
24    KSP,
25    SNES,
26    TS,
27    TAO,
28    TAOLineSearch,
29    DM,
30)
31
32__all__ = [
33    'Scalar',
34    'ArrayBool',
35    'ArrayInt',
36    'ArrayReal',
37    'ArrayComplex',
38    'ArrayScalar',
39    'DimsSpec',
40    'AccessModeSpec',
41    'InsertModeSpec',
42    'ScatterModeSpec',
43    'LayoutSizeSpec',
44    'NormTypeSpec',
45    'PetscOptionsHandlerFunction',
46    'MatAssemblySpec',
47    'MatSizeSpec',
48    'MatBlockSizeSpec',
49    'CSRIndicesSpec',
50    'CSRSpec',
51    'NNZSpec',
52    'MatNullFunction',
53    'DMCoarsenHookFunction',
54    'DMRestrictHookFunction',
55    'KSPRHSFunction',
56    'KSPOperatorsFunction',
57    'KSPConvergenceTestFunction',
58    'KSPMonitorFunction',
59    'KSPPreSolveFunction',
60    'KSPPostSolveFunction',
61    'SNESMonitorFunction',
62    'SNESObjFunction',
63    'SNESFunction',
64    'SNESJacobianFunction',
65    'SNESGuessFunction',
66    'SNESUpdateFunction',
67    'SNESLSPreFunction',
68    'SNESNGSFunction',
69    'SNESConvergedFunction',
70    'TSRHSFunction',
71    'TSRHSJacobian',
72    'TSRHSJacobianP',
73    'TSIFunction',
74    'TSIJacobian',
75    'TSIJacobianP',
76    'TSI2Function',
77    'TSI2Jacobian',
78    'TSI2JacobianP',
79    'TSMonitorFunction',
80    'TSPreStepFunction',
81    'TSPostStepFunction',
82    'TSIndicatorFunction',
83    'TSPostEventFunction',
84    'TSPreStepFunction',
85    'TSPostStepFunction',
86    'TAOObjectiveFunction',
87    'TAOGradientFunction',
88    'TAOObjectiveGradientFunction',
89    'TAOHessianFunction',
90    'TAOUpdateFunction',
91    'TAOMonitorFunction',
92    'TAOConvergedFunction',
93    'TAOJacobianFunction',
94    'TAOResidualFunction',
95    'TAOJacobianResidualFunction',
96    'TAOVariableBoundsFunction',
97    'TAOConstraintsFunction',
98    'TAOConstraintsJacobianFunction',
99    'TAOLSObjectiveFunction',
100    'TAOLSGradientFunction',
101    'TAOLSObjectiveGradientFunction',
102]
103
104# --- Sys ---
105
106Scalar = float | complex
107"""Scalar type.
108
109Scalars can be either `float` or `complex` (but not both) depending on how
110PETSc was configured (``./configure --with-scalar-type=real|complex``).
111
112"""
113
114ArrayBool = NDArray[np.bool_]
115"""Array of `bool`."""
116
117ArrayInt = NDArray[np.integer]
118"""Array of `int`."""
119
120ArrayReal = NDArray[np.floating]
121"""Array of `float`."""
122
123ArrayComplex = NDArray[np.complexfloating]
124"""Array of `complex`."""
125
126ArrayScalar = NDArray[np.floating | np.complexfloating]
127"""Array of `Scalar` numbers."""
128
129DimsSpec = tuple[int, ...]
130"""Dimensions specification.
131
132   N-tuples indicates N-dimensional grid sizes.
133
134"""
135
136AccessModeSpec = Literal['rw', 'r', 'w'] | None
137"""Access mode specification.
138
139   Possible values are:
140     - ``'rw'`` Read-Write mode.
141     - ``'r'`` Read-only mode.
142     - ``'w'`` Write-only mode.
143     - `None` as ``'rw'``.
144
145"""
146
147InsertModeSpec: TypeAlias = InsertMode | bool | None
148"""Insertion mode specification.
149
150   Possible values are:
151     - `InsertMode.ADD_VALUES` Add new value to existing one.
152     - `InsertMode.INSERT_VALUES` Replace existing entry with new value.
153     - `None` as `InsertMode.INSERT_VALUES`.
154     - `False` as `InsertMode.INSERT_VALUES`.
155     - `True` as `InsertMode.ADD_VALUES`.
156
157   See Also
158   --------
159   InsertMode
160
161"""
162
163ScatterModeSpec: TypeAlias = ScatterMode | bool | str | None
164"""Scatter mode specification.
165
166   Possible values are:
167     - `ScatterMode.FORWARD` Forward mode.
168     - `ScatterMode.REVERSE` Reverse mode.
169     - `None` as `ScatterMode.FORWARD`.
170     - `False` as `ScatterMode.FORWARD`.
171     - `True` as `ScatterMode.REVERSE`.
172     - ``'forward'`` as `ScatterMode.FORWARD`.
173     - ``'reverse'`` as `ScatterMode.REVERSE`.
174
175   See Also
176   --------
177   ScatterMode
178
179"""
180
181LayoutSizeSpec: TypeAlias = int | tuple[int, int]
182"""`int` or 2-`tuple` of `int` describing the layout sizes.
183
184   A single `int` indicates global size.
185   A `tuple` of `int` indicates ``(local_size, global_size)``.
186
187   See Also
188   --------
189   Sys.splitOwnership
190
191"""
192
193NormTypeSpec: TypeAlias = NormType | None
194"""Norm type specification.
195
196    Possible values include:
197
198    - `NormType.NORM_1` The 1-norm: Σₙ abs(xₙ) for vectors, maxₙ (Σᵢ abs(xₙᵢ)) for matrices.
199    - `NormType.NORM_2` The 2-norm: √(Σₙ xₙ²) for vectors, largest singular values for matrices.
200    - `NormType.NORM_INFINITY` The ∞-norm: maxₙ abs(xₙ) for vectors, maxᵢ (Σₙ abs(xₙᵢ)) for matrices.
201    - `NormType.NORM_FROBENIUS` The Frobenius norm: same as 2-norm for vectors, √(Σₙᵢ xₙᵢ²) for matrices.
202    - `NormType.NORM_1_AND_2` Compute both `NormType.NORM_1` and `NormType.NORM_2`.
203    - `None` as `NormType.NORM_2` for vectors, `NormType.NORM_FROBENIUS` for matrices.
204
205    See Also
206    --------
207    PETSc.NormType, petsc.NormType
208
209"""
210
211# --- PetscObject ---
212
213PetscOptionsHandlerFunction: TypeAlias = Callable[[Object], None]
214"""Callback for processing extra options."""
215
216# --- Mat ---
217
218MatAssemblySpec: TypeAlias = Mat.AssemblyType | bool | None
219"""Matrix assembly specification.
220
221   Possible values are:
222     - `Mat.AssemblyType.FINAL`
223     - `Mat.AssemblyType.FLUSH`
224     - `None` as `Mat.AssemblyType.FINAL`
225     - `False` as `Mat.AssemblyType.FINAL`
226     - `True` as `Mat.AssemblyType.FLUSH`
227
228   See Also
229   --------
230   petsc.MatAssemblyType
231
232"""
233
234MatSizeSpec: TypeAlias = int | tuple[int, int] | tuple[tuple[int, int], tuple[int, int]]
235"""`int` or (nested) `tuple` of `int` describing the matrix sizes.
236
237   If `int` then rows = columns.
238   A single `tuple` of `int` indicates ``(rows, columns)``.
239   A nested `tuple` of `int` indicates ``((local_rows, rows), (local_columns, columns))``.
240
241   See Also
242   --------
243   Sys.splitOwnership
244
245"""
246
247MatBlockSizeSpec: TypeAlias = int | tuple[int, int]
248"""The row and column block sizes.
249
250   If a single `int` is provided then rows and columns share the same block size.
251
252"""
253
254CSRIndicesSpec: TypeAlias = tuple[Sequence[int], Sequence[int]]
255"""CSR indices format specification.
256
257   A 2-tuple carrying the ``(row_start, col_indices)`` information.
258
259"""
260
261CSRSpec: TypeAlias = tuple[Sequence[int], Sequence[int], Sequence[Scalar]]
262"""CSR format specification.
263
264   A 3-tuple carrying the ``(row_start, col_indices, values)`` information.
265
266"""
267
268NNZSpec: TypeAlias = int | Sequence[int] | tuple[Sequence[int], Sequence[int]]
269"""Nonzero pattern specification.
270
271   A single `int` corresponds to fixed number of non-zeros per row.
272   A `Sequence` of `int` indicates different non-zeros per row.
273   If a 2-`tuple` is used, the elements of the tuple corresponds
274   to the on-process and off-process parts of the matrix.
275
276   See Also
277   --------
278   petsc.MatSeqAIJSetPreallocation, petsc.MatMPIAIJSetPreallocation
279
280"""
281
282# --- MatNullSpace ---
283
284MatNullFunction = Callable[[NullSpace, Vec], None]
285"""`PETSc.NullSpace` callback."""
286
287# --- DM ---
288
289DMCoarsenHookFunction = Callable[[DM, DM], None]
290"""`PETSc.DM` coarsening hook callback."""
291
292DMRestrictHookFunction = Callable[[DM, Mat, Vec, Mat, DM], None]
293"""`PETSc.DM` restriction hook callback."""
294
295# --- KSP ---
296
297KSPRHSFunction = Callable[[KSP, Vec], None]
298"""`PETSc.KSP` right-hand side function callback."""
299
300KSPOperatorsFunction = Callable[[KSP, Mat, Mat], None]
301"""`PETSc.KSP` operators function callback."""
302
303KSPConvergenceTestFunction = Callable[[KSP, int, float], KSP.ConvergedReason]
304"""`PETSc.KSP` convergence test callback."""
305
306KSPMonitorFunction = Callable[[KSP, int, float], None]
307"""`PETSc.KSP` monitor callback."""
308
309KSPPreSolveFunction = Callable[[KSP, Vec, Vec], None]
310"""`PETSc.KSP` pre solve callback."""
311
312KSPPostSolveFunction = Callable[[KSP, Vec, Vec], None]
313"""`PETSc.KSP` post solve callback."""
314
315# --- SNES ---
316
317SNESMonitorFunction = Callable[[SNES, int, float], None]
318"""`SNES` monitor callback."""
319
320SNESObjFunction = Callable[[SNES, Vec], None]
321"""`SNES` objective function callback."""
322
323SNESFunction = Callable[[SNES, Vec, Vec], None]
324"""`SNES` residual function callback."""
325
326SNESJacobianFunction = Callable[[SNES, Vec, Mat, Mat], None]
327"""`SNES` Jacobian callback."""
328
329SNESGuessFunction = Callable[[SNES, Vec], None]
330"""`SNES` initial guess callback."""
331
332SNESUpdateFunction = Callable[[SNES, int], None]
333"""`SNES` step update callback."""
334
335SNESLSPreFunction = Callable[[Vec, Vec], None]
336"""`SNES` linesearch pre-check update callback."""
337
338SNESNGSFunction = Callable[[SNES, Vec, Vec], None]
339"""`SNES` nonlinear Gauss-Seidel callback."""
340
341SNESConvergedFunction = Callable[
342    [SNES, int, tuple[float, float, float]], SNES.ConvergedReason
343]
344"""`SNES` convergence test callback."""
345
346# --- TS ---
347
348TSRHSFunction = Callable[[TS, float, Vec, Vec], None]
349"""`TS` right-hand side function callback."""
350
351TSRHSJacobian = Callable[[TS, float, Vec, Mat, Mat], None]
352"""`TS` right-hand side Jacobian callback."""
353
354TSRHSJacobianP = Callable[[TS, float, Vec, Mat], None]
355"""`TS` right-hand side parameter Jacobian callback."""
356
357TSIFunction = Callable[[TS, float, Vec, Vec, Vec], None]
358"""`TS` implicit function callback."""
359
360TSIJacobian = Callable[[TS, float, Vec, Vec, float, Mat, Mat], None]
361"""`TS` implicit Jacobian callback."""
362
363TSIJacobianP = Callable[[TS, float, Vec, Vec, float, Mat], None]
364"""`TS` implicit parameter Jacobian callback."""
365
366TSI2Function = Callable[[TS, float, Vec, Vec, Vec, Vec], None]
367"""`TS` implicit 2nd order function callback."""
368
369TSI2Jacobian = Callable[[TS, float, Vec, Vec, Vec, float, float, Mat, Mat], None]
370"""`TS` implicit 2nd order Jacobian callback."""
371
372TSI2JacobianP = Callable[[TS, float, Vec, Vec, Vec, float, float, Mat], None]
373"""`TS` implicit 2nd order parameter Jacobian callback."""
374
375TSMonitorFunction = Callable[[TS, int, float, Vec], None]
376"""`TS` monitor callback."""
377
378TSPreStepFunction = Callable[[TS], None]
379"""`TS` pre-step callback."""
380
381TSPostStepFunction = Callable[[TS], None]
382"""`TS` post-step callback."""
383
384TSIndicatorFunction = Callable[[TS, float, Vec, NDArray[np.floating]], None]
385"""`TS` event indicator callback."""
386
387TSPostEventFunction = Callable[[TS, NDArray[np.integer], float, Vec, bool], None]
388"""`TS` post-event callback."""
389
390# --- TAO ---
391
392TAOObjectiveFunction = Callable[[TAO, Vec], float]
393"""`TAO` objective function callback."""
394
395TAOGradientFunction = Callable[[TAO, Vec, Vec], None]
396"""`TAO` objective gradient callback."""
397
398TAOObjectiveGradientFunction = Callable[[TAO, Vec, Vec], float]
399"""`TAO` objective function and gradient callback."""
400
401TAOHessianFunction = Callable[[TAO, Vec, Mat, Mat], None]
402"""`TAO` objective Hessian callback."""
403
404TAOUpdateFunction = Callable[[TAO, int], None]
405"""`TAO` update callback."""
406
407TAOMonitorFunction = Callable[[TAO], None]
408"""`TAO` monitor callback."""
409
410TAOConvergedFunction = Callable[[TAO], None]
411"""`TAO` convergence test callback."""
412
413TAOJacobianFunction = Callable[[TAO, Vec, Mat, Mat], None]
414"""`TAO` Jacobian callback."""
415
416TAOResidualFunction = Callable[[TAO, Vec, Vec], None]
417"""`TAO` residual callback."""
418
419TAOJacobianResidualFunction = Callable[[TAO, Vec, Mat, Mat], None]
420"""`TAO` Jacobian residual callback."""
421
422TAOVariableBoundsFunction = Callable[[TAO, Vec, Vec], None]
423"""`TAO` variable bounds callback."""
424
425TAOConstraintsFunction = Callable[[TAO, Vec, Vec], None]
426"""`TAO` constraints callback."""
427
428TAOConstraintsJacobianFunction = Callable[[TAO, Vec, Mat, Mat], None]
429"""`TAO` constraints Jacobian callback."""
430
431TAOLSObjectiveFunction = Callable[[TAOLineSearch, Vec], float]
432"""`TAOLineSearch` objective function callback."""
433
434TAOLSGradientFunction = Callable[[TAOLineSearch, Vec, Vec], None]
435"""`TAOLineSearch` objective gradient callback."""
436
437TAOLSObjectiveGradientFunction = Callable[[TAOLineSearch, Vec, Vec], float]
438"""`TAOLineSearch` objective function and gradient callback."""
439