Coverage for pyrc\core\solver\symbolic.py: 100%

1# ------------------------------------------------------------------------------- 

2# Copyright (C) 2026 Joel Kimmich, Tim Jourdan 

3# ------------------------------------------------------------------------------ 

4# License 

5# This file is part of PyRC, distributed under GPL-3.0-or-later. 

6# ------------------------------------------------------------------------------ 

7 

8from abc import ABC, abstractmethod 

9from typing import Any 

10 

11import numpy as np 

12from scipy.sparse import coo_matrix, csr_matrix 

13from sympy import SparseMatrix, Symbol, lambdify 

14 

15 

16class SymbolicEvaluator(ABC): 

17 def __init__(self, entries: dict[tuple[int, int], Any], time_symbols: list[Symbol]): 

18 """ 

19 Parameters 

20 ---------- 

21 entries : dict[tuple[int, int], any] 

22 Mapping from (i, j) index to sympy expression or numeric value. 

23 time_symbols : list[Symbol] 

24 Ordered list of time-dependent symbols. 

25 """ 

26 self.time_symbols: list[Symbol] = time_symbols 

27 

28 symbolic_expressions, symbolic_indices = [], [] 

29 constant_values, constant_indices = [], [] 

30 

31 for index, expr in entries.items(): 

32 if hasattr(expr, "free_symbols") and expr.free_symbols: 

33 symbolic_expressions.append(expr) 

34 symbolic_indices.append(index) 

35 else: 

36 constant_values.append(float(expr)) 

37 constant_indices.append(index) 

38 

39 self._n_symbolic = len(symbolic_expressions) 

40 n_constant = len(constant_values) 

41 

42 if self._n_symbolic > 0: 

43 self._symbolic_func = lambdify(time_symbols, symbolic_expressions, "numpy") 

44 

45 self._all_data = np.empty(n_constant + self._n_symbolic, dtype=np.float64) 

46 self._all_data[:n_constant] = constant_values 

47 self._symbolic_slice = slice(n_constant, n_constant + self._n_symbolic) 

48 

49 all_indices = constant_indices + symbolic_indices 

50 self._build_structure(all_indices, self._all_data) 

51 

52 @abstractmethod 

53 def _build_structure(self, all_indices: list[tuple[int, int]], all_data: np.ndarray) -> None: 

54 """Build the output structure from parsed indices and data array. 

55 

56 Parameters 

57 ---------- 

58 all_indices : list[tuple[int, int]] 

59 All (i, j) indices in constant-first order. 

60 all_data : np.ndarray 

61 Pre-allocated data array; constants already filled. 

62 """ 

63 

64 @abstractmethod 

65 def _scatter(self) -> None: 

66 """Scatter self._all_data into the output structure in-place.""" 

67 

68 def evaluate(self, time_values=None): 

69 """ 

70 Parameters 

71 ---------- 

72 time_values : array-like, optional 

73 Values for time-dependent symbols in same order as time_symbols. 

74 

75 Returns 

76 ------- 

77 Output structure with evaluated values. 

78 """ 

79 if self._n_symbolic > 0: 

80 self._all_data[self._symbolic_slice] = self._symbolic_func(*time_values) 

81 self._scatter() 

82 return self._output 

83 

84 

85class SparseSymbolicEvaluator(SymbolicEvaluator): 

86 def __init__(self, semi_symbolic_matrix: SparseMatrix, time_symbols: list[Symbol]): 

87 """ 

88 Parameters 

89 ---------- 

90 semi_symbolic_matrix : SparseMatrix 

91 Sparse matrix with time-dependent symbols or numeric values. 

92 time_symbols : list[Symbol] 

93 Ordered list of time-dependent symbols. 

94 """ 

95 self.shape = semi_symbolic_matrix.shape 

96 entries = {(i, j): semi_symbolic_matrix[i, j] for (i, j) in semi_symbolic_matrix.todok()} 

97 super().__init__(entries, time_symbols) 

98 

99 def _build_structure(self, all_indices: list[tuple[int, int]], all_data: np.ndarray) -> None: 

100 all_rows = np.array([i for i, _ in all_indices], dtype=np.int32) 

101 all_cols = np.array([j for _, j in all_indices], dtype=np.int32) 

102 self._perm = np.lexsort((all_cols, all_rows)) 

103 coo = coo_matrix((all_data, (all_rows, all_cols)), shape=self.shape) 

104 self._output: csr_matrix = csr_matrix(coo) 

105 self._csr_data = self._output.data 

106 

107 def _scatter(self) -> None: 

108 self._csr_data[:] = self._all_data[self._perm] 

109 

110 

111class ArraySymbolicEvaluator(SymbolicEvaluator): 

112 def __init__(self, semi_symbolic_array: np.ndarray, time_symbols: list[Symbol]): 

113 """ 

114 Parameters 

115 ---------- 

116 semi_symbolic_array : np.ndarray 

117 1D or column array (shape (n,) or (n, 1)) with sympy expressions or numeric values. 

118 time_symbols : list[Symbol] 

119 Ordered list of time-dependent symbols. 

120 """ 

121 self._original_shape = semi_symbolic_array.shape 

122 flat = semi_symbolic_array.ravel() 

123 entries = {(i, 0): expr for i, expr in enumerate(flat)} 

124 super().__init__(entries, time_symbols) 

125 

126 def _build_structure(self, all_indices: list[tuple[int, int]], all_data: np.ndarray) -> None: 

127 self._scatter_indices = np.array([i for i, _ in all_indices], dtype=np.int32) 

128 self._flat_output = np.empty(len(all_indices), dtype=np.float64) 

129 self._flat_output[self._scatter_indices] = all_data 

130 self._output = self._flat_output.reshape(self._original_shape) 

131 

132 def _scatter(self) -> None: 

133 self._flat_output[self._scatter_indices] = self._all_data