Coverage for src/amisc/compression.py: 91%
139 statements
« prev ^ index » next coverage.py v7.6.1, created at 2024-11-05 19:26 +0000
1"""Module for compression methods.
3Especially useful for field quantities with high dimensions.
5Includes:
7- `Compression` — an interface for specifying a compression method for field quantities.
8- `SVD` — a Singular Value Decomposition (SVD) compression method.
9"""
10from __future__ import annotations
12from abc import ABC, abstractmethod
13from dataclasses import dataclass, field
15import numpy as np
16from scipy.interpolate import RBFInterpolator
18from amisc.serialize import PickleSerializable
20__all__ = ["Compression", "SVD"]
23@dataclass
24class Compression(PickleSerializable, ABC):
25 """Base class for compression methods. Compression methods should:
27 - `compute_map` - compute the compression map from provided data
28 - `compress` - compress data into a latent space
29 - `reconstruct` - reconstruct the compressed data back into the full space
30 - `latent_size` - return the size of the latent space
31 - `estimate_latent_ranges` - estimate the range of the latent space coefficients
33 !!! Note "Specifying fields"
34 The `fields` attribute is a list of strings that specify the field quantities to compress. For example, for
35 3D velocity data, the fields might be `['ux', 'uy', 'uz']`. The length of the
36 `fields` attribute is used to determine the number of quantities of interest at each grid point in `coords`.
37 Note that interpolation to/from the compression grid will assume a shape of `(num_pts, num_qoi)` for the
38 states on the grid, where `num_qoi` is the length of `fields` and `num_pts` is the length of `coords`. When
39 constructing the compression map, this important fact should be considered when passing data to
40 `compute_map`.
42 In order to use a `Compression` object, you must first call `compute_map` to compute the compression map, which
43 should set the private value `self._map_computed=True`. The `coords` of the compression grid must also be
44 specified. The `coords` should have the shape `(num_pts, dim)` where `num_pts` is the number of points in the
45 compression grid and `dim` is the number of spatial dimensions. If `coords` is a 1d array, then the `dim` is
46 assumed to be 1.
48 :ivar fields: list of field quantities to compress
49 :ivar method: the compression method to use (only svd is supported for now)
50 :ivar coords: the coordinates of the compression grid
51 :ivar interpolate_method: the interpolation method to use to interpolate to/from the compression grid
52 (only `rbf` (i.e. radial basis function) is supported for now)
53 :ivar interpolate_opts: additional options to pass to the interpolation method
54 :ivar _map_computed: whether the compression map has been computed
55 """
56 fields: list[str] = field(default_factory=list)
57 method: str = 'svd'
58 coords: np.ndarray = None # (num_pts, dim)
59 interpolate_method: str = 'rbf'
60 interpolate_opts: dict = field(default_factory=dict)
61 _map_computed: bool = False
63 @property
64 def map_exists(self):
65 """All compression methods should have `coords` when their map has been constructed."""
66 return self.coords is not None and self._map_computed
68 @property
69 def dim(self):
70 """Number of physical grid coordinates for the field quantity, (i.e. x,y,z spatial dims)"""
71 return self.coords.shape[1] if (self.coords is not None and len(self.coords.shape) > 1) else 1
73 @property
74 def num_pts(self):
75 """Number of physical points in the compression grid."""
76 return self.coords.shape[0] if self.coords is not None else None
78 @property
79 def num_qoi(self):
80 """Number of quantities of interest at each grid point, (i.e. `ux, uy, uz` for 3d velocity data)."""
81 return len(self.fields) if self.fields is not None else 1
83 @property
84 def dof(self):
85 """Total degrees of freedom in the compression grid (i.e. `num_pts * num_qoi`)."""
86 return self.num_pts * self.num_qoi if self.num_pts is not None else None
88 def _correct_coords(self, coords):
89 """Correct the coordinates to be in the correct shape for compression."""
90 coords = np.atleast_1d(coords)
91 if len(coords.shape) == 1:
92 coords = coords[..., np.newaxis] if self.dim == 1 else coords[np.newaxis, ...]
93 return coords
95 def interpolator(self):
96 """The interpolator to use during compression and reconstruction. Interpolator expects to be used as:
98 ```python
99 xg = np.ndarray # (num_pts, dim) grid coordinates
100 yg = np.ndarray # (num_pts, ...) scalar values on grid
101 xp = np.ndarray # (Q, dim) evaluation points
103 interp = interpolate_method(xg, yg, **interpolate_opts)
105 yp = interp(xp) # (Q, ...) interpolated values
106 ```
107 """
108 method = self.interpolate_method or 'rbf'
109 match method.lower():
110 case 'rbf':
111 return RBFInterpolator
112 case other:
113 raise NotImplementedError(f"Interpolation method '{other}' is not implemented.")
115 def interpolate_from_grid(self, states: np.ndarray, new_coords: np.ndarray):
116 """Interpolate the states on the compression grid to new coordinates.
118 :param states: `(*loop_shape, dof)` - the states on the compression grid
119 :param new_coords: `(*coord_shape, dim)` - the new coordinates to interpolate to
120 :return: `dict` of `(*loop_shape, *coord_shape)` for each qoi - the interpolated states
121 """
122 new_coords = self._correct_coords(new_coords)
123 grid_coords = self._correct_coords(self.coords)
124 skip_interp = (new_coords.shape == grid_coords.shape and np.allclose(new_coords, grid_coords))
126 ret_dict = {}
127 loop_shape = states.shape[:-1]
128 coords_shape = new_coords.shape[:-1]
129 states = states.reshape((*loop_shape, self.num_pts, self.num_qoi))
130 new_coords = new_coords.reshape((-1, self.dim))
131 for i, qoi in enumerate(self.fields):
132 if skip_interp:
133 ret_dict[qoi] = states[..., i]
134 else:
135 reshaped_states = states[..., i].reshape(-1, self.num_pts).T # (num_pts, ...)
136 interp = self.interpolator()(grid_coords, reshaped_states, **self.interpolate_opts)
137 yp = interp(new_coords)
138 ret_dict[qoi] = yp.T.reshape(*loop_shape, *coords_shape)
140 return ret_dict
142 def interpolate_to_grid(self, field_coords: np.ndarray, field_values):
143 """Interpolate the field values at given coordinates to the compression grid.
145 :param field_coords: `(*coord_shape, dim)` - the coordinates of the field values
146 :param field_values: `dict` of `(*loop_shape, *coord_shape)` for each qoi - the field values at the coordinates
147 :return: `(*loop_shape, dof)` - the interpolated values on the compression grid
148 """
149 field_coords = self._correct_coords(field_coords)
150 grid_coords = self._correct_coords(self.coords)
151 skip_interp = (field_coords.shape == grid_coords.shape and np.allclose(field_coords, grid_coords))
153 coords_shape = field_coords.shape[:-1]
154 loop_shape = next(iter(field_values.values())).shape[:-len(coords_shape)]
155 states = np.empty((*loop_shape, self.num_pts, self.num_qoi))
156 field_coords = field_coords.reshape(-1, self.dim)
157 for i, qoi in enumerate(self.fields):
158 field_vals = field_values[qoi].reshape((*loop_shape, -1)) # (..., Q)
159 if skip_interp:
160 states[..., i] = field_vals
161 else:
162 field_vals = field_vals.reshape((-1, field_vals.shape[-1])).T # (Q, ...)
163 interp = self.interpolator()(field_coords, field_vals, **self.interpolate_opts)
164 yg = interp(grid_coords)
165 states[..., i] = yg.T.reshape(*loop_shape, self.num_pts)
167 return states.reshape((*loop_shape, self.dof))
169 @abstractmethod
170 def compute_map(self, **kwargs):
171 """Compute and store the compression map. Must set the value of `coords` and `_is_computed`. Should
172 use the same normalization as the parent `Variable` object.
174 !!! Note
175 You should pass any required data to `compute_map` with the assumption that the data will be used in the
176 shape `(num_pts, num_qoi)` where `num_qoi` is the length of `fields` and `num_pts` is the length of
177 `coords`. This is the shape that the compression map should be constructed in.
178 """
179 raise NotImplementedError
181 @abstractmethod
182 def compress(self, data: np.ndarray) -> np.ndarray:
183 """Compress the data into a latent space.
185 :param data: `(..., dof)` - the data to compress from full size of `dof`
186 :return: `(..., rank)` - the compressed latent space data with size `rank`
187 """
188 raise NotImplementedError
190 @abstractmethod
191 def reconstruct(self, compressed: np.ndarray) -> np.ndarray:
192 """Reconstruct the compressed data back into the full `dof` space.
194 :param compressed: `(..., rank)` - the compressed data to reconstruct
195 :return: `(..., dof)` - the reconstructed data with full `dof`
196 """
197 raise NotImplementedError
199 @abstractmethod
200 def latent_size(self) -> int:
201 """Return the size of the latent space."""
202 raise NotImplementedError
204 @abstractmethod
205 def estimate_latent_ranges(self) -> list[tuple[float, float]]:
206 """Estimate the range of the latent space coefficients."""
207 raise NotImplementedError
209 @classmethod
210 def from_dict(cls, spec: dict) -> Compression:
211 """Construct a `Compression` object from a spec dictionary."""
212 method = spec.pop('method', 'svd').lower()
213 match method:
214 case 'svd':
215 return SVD(**spec)
216 case other:
217 raise NotImplementedError(f"Compression method '{other}' is not implemented.")
220@dataclass
221class SVD(Compression):
222 """A Singular Value Decomposition (SVD) compression method. The SVD will be computed on initialization if the
223 `data_matrix` is provided.
225 :ivar data_matrix: `(dof, num_samples)` - the data matrix
226 :ivar projection_matrix: `(dof, rank)` - the projection matrix
227 :ivar rank: the rank of the SVD decomposition
228 :ivar energy_tol: the energy tolerance of the SVD decomposition
229 """
230 data_matrix: np.ndarray = None # (dof, num_samples)
231 projection_matrix: np.ndarray = None # (dof, rank)
232 rank: int = None
233 energy_tol: float = None
235 def __post_init__(self):
236 """Compute the SVD if the data matrix is provided."""
237 if (data_matrix := self.data_matrix) is not None:
238 self.compute_map(data_matrix, rank=self.rank, energy_tol=self.energy_tol)
240 def compute_map(self, data_matrix: np.ndarray | dict, rank: int = None, energy_tol: float = None):
241 """Compute the SVD compression map from the data matrix. Recall that `dof` is the total number of degrees of
242 freedom, equal to the number of grid points `num_pts` times the number of quantities of interest `num_qoi`
243 at each grid point.
245 :param data_matrix: `(dof, num_samples)` - the data matrix. If passed in as a `dict`, then the data matrix
246 will be formed by concatenating the values of the `dict` along the last axis in the order
247 of the `fields` attribute and flattening the last two axes. This is useful for passing
248 in a dictionary of field values like `{field1: (num_samples, num_pts), field2: ...}`
249 which ensures consistency of shape with the compression `coords`.
250 :param rank: the rank of the SVD decomposition
251 :param energy_tol: the energy tolerance of the SVD decomposition
252 """
253 if isinstance(data_matrix, dict):
254 data_matrix = np.concatenate([data_matrix[field][..., np.newaxis] for field in self.fields], axis=-1)
255 data_matrix = data_matrix.reshape(*data_matrix.shape[:-2], -1).T # (dof, num_samples)
257 nan_idx = np.any(np.isnan(data_matrix), axis=0)
258 data_matrix = data_matrix[:, ~nan_idx]
259 u, s, vt = np.linalg.svd(data_matrix)
260 energy_frac = np.cumsum(s ** 2 / np.sum(s ** 2))
261 if rank := (rank or self.rank):
262 energy_tol = energy_frac[rank - 1]
263 else:
264 energy_tol = energy_tol or self.energy_tol or 0.95
265 idx = int(np.where(energy_frac >= energy_tol)[0][0])
266 rank = idx + 1
268 self.data_matrix = data_matrix
269 self.rank = rank
270 self.energy_tol = energy_tol
271 self.projection_matrix = u[:, :rank] # (dof, rank)
272 self._map_computed = True
274 def compress(self, data):
275 return np.squeeze(self.projection_matrix.T @ data[..., np.newaxis], axis=-1)
277 def reconstruct(self, compressed):
278 return np.squeeze(self.projection_matrix @ compressed[..., np.newaxis], axis=-1)
280 def latent_size(self):
281 return self.rank
283 def estimate_latent_ranges(self):
284 if self.map_exists:
285 latent_data = self.compress(self.data_matrix.T) # (rank, num_samples)
286 latent_min = np.min(latent_data, axis=0)
287 latent_max = np.max(latent_data, axis=0)
288 return [(lmin, lmax) for lmin, lmax in zip(latent_min, latent_max)]
289 else:
290 return None