Skip to content

Optimisation #12

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
perrymok wants to merge 11 commits into
base: main
Choose a base branch
from
Open

Optimisation #12

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
11 commits
Select commit Hold shift + click to select a range
b7e2d53
vectorisation and list/dict comprehension
perrymok Jul 22, 2026
2b9e579
use `np.vdot(A, B)` and `np.outer(A, B)` instead of `np.sum(A * B)` a…
perrymok Jul 22, 2026
7f08b17
caching to reduce disk usage
perrymok Jul 22, 2026
2e52dff
only assert in safe mode
perrymok Jul 22, 2026
6c5e939
remove redundant clipping (ATTENTION REQUIRED)
perrymok Jul 22, 2026
8586ac6
use more efficient operations
perrymok Jul 22, 2026
b797c11
revert 2e52dff4
perrymok Jul 22, 2026
2b5da69
Merge remote-tracking branch 'origin/main' into optimisation
perrymok Jul 29, 2026
f4dfb6d
associate to matrix multiplication with lower FLOPs
perrymok Jul 30, 2026
3ce4803
only recompute changed parts
perrymok Jul 30, 2026
2e541ce
Merge remote-tracking branch 'origin/main' into optimisation
perrymok Aug 5, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
170 changes: 68 additions & 102 deletions espm/datasets/eds_spim.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -10,7 +10,6 @@
- set fixed W for the :class:`espm.estimators.NMFEstimator` decomposition
"""

import json
import warnings

import hyperspy.api as hs
Expand All @@ -31,21 +30,19 @@
from scipy.optimize import curve_fit
from tqdm import tqdm

from espm.conf import NUMBER_PERIODIC_TABLE, SYMBOLS_PERIODIC_TABLE
from espm.estimators import NMFEstimator, SmoothNMF
from espm.models import EDXS
from espm.utils import (
get_explained_intensity_W,
get_npt,
get_spt,
num_to_symbol,
number_to_symbol_list,
quant_spectrum,
symbol_to_number_dict,
symbol_to_number_list,
)

# TODO: use cache after merging optimisation
NPT = json.load(open(NUMBER_PERIODIC_TABLE))


def _check_decomposition(func) -> bool:
# @wraps(func)
Expand Down Expand Up @@ -407,15 +404,12 @@ def carto_fixed_W(self, brstlg_comps: int = 1) -> np.ndarray:
)
elements = self.metadata.EDS_model.elements
if self.problem_type == "no_brstlg":
W = np.diag(-1 * np.ones((len(elements),)))
W = np.diag(np.full((len(elements),), -1.0))
elif self.problem_type == "bremsstrahlung":
W1 = np.diag(-1 * np.ones((len(elements),)))
W2 = np.zeros((2, len(elements)))
W_elts = np.vstack((W1, W2))
W3 = np.zeros((len(elements), brstlg_comps))
W4 = -1 * np.ones((2, brstlg_comps))
W_brstlg = np.vstack((W3, W4))
W = np.hstack((W_elts, W_brstlg))
L = len(elements)

@perrymok perrymok Jul 30, 2026 •
edited
Loading

Copy link
Copy Markdown
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This is quite a big change but if u look closely it's way more elegant than the original method, and actually much easier to understand

W = np.zeros((L + 2, L + brstlg_comps))
np.fill_diagonal(W[:L, :L], -1.0)
W[L:, L:] = -1.0

return W

Expand Down Expand Up @@ -449,31 +443,26 @@ def convert_to_symbols(elements=[]):
conv_elts = convert_to_symbols(elements=elements)

if self.problem_type == "no_brstlg":
W = -1 * np.ones((len(raw_elts), len(phases_dict.keys())))
W = np.full((len(raw_elts), len(phases_dict.keys())), -1.0)
elif self.problem_type == "bremsstrahlung":
W = -1 * np.ones((len(raw_elts) + 2, len(phases_dict.keys())))
W = np.full((len(raw_elts) + 2, len(phases_dict.keys())), -1.0)
else:
raise ValueError(
"problem type should be either no_brstlg or bremsstrahlung"
)
for p, phase in enumerate(phases_dict):
for key in phases_dict[phase]:
if key == "b0":
if self.problem_type == "bremsstrahlung":
W[-2, p] = phases_dict[phase][key]
else:
warnings.warn(
"The chosen EDXS modelling does not incorporate the bremsstrahlung. Input bremsstrahlung parameters will be ignored."
)
if key == "b1":
if self.problem_type == "bremsstrahlung":
W[-1, p] = phases_dict[phase][key]
else:
warnings.warn(
"The chosen EDXS modelling does not incorporate the bremsstrahlung. Input bremsstrahlung parameters will be ignored."
)
if key in conv_elts:
W[indices[conv_elts.index(key)], p] = phases_dict[phase][key]
match key:
case "b0" | "b1":
if self.problem_type == "bremsstrahlung":
idx = -2 if key == "b0" else -1
W[idx, p] = phases_dict[phase][key]
else:
warnings.warn(
"The chosen EDXS modelling does not incorporate the bremsstrahlung. Input bremsstrahlung parameters will be ignored."
)
case _ if key in conv_elts:
W[indices[conv_elts.index(key)], p] = phases_dict[phase][key]
return W

def print_concentration_report(
Expand Down Expand Up @@ -591,11 +580,11 @@ def estimate_mass_thickness(
W_init = estimator.W_
elts = list(self.model.get_elements(include_ignored=False))
elts_indices = self.model.NMF_simplex()
new_elts_dict = {elts[i]: W_init[elts_indices[i]] for i in range(len(elts))}
new_elts_dict = dict(zip(elts, W_init[elts_indices]))

_ = 0
curr_mt = self.metadata.Sample.thickness * self.metadata.Sample.density
while _ < 5:

for _ in range(5):
# first init of the model
brstlg_model, mask = self.model.bremsstrahlung_only_tools(
mass_thickness=curr_mt, elements_dict=new_elts_dict, ranges=self.ranges
Expand All @@ -607,12 +596,12 @@ def estimate_mass_thickness(
brstlg_estimator.fit(masked_X[:, np.newaxis])
W_brstlg = np.vstack(
(
-1
* np.ones(
np.full(
(
W_init.shape[0] - brstlg_estimator.W_.shape[0],
brstlg_estimator.W_.shape[1],
)
),
-1.0,
),
brstlg_estimator.W_,
)
Expand All @@ -629,14 +618,10 @@ def estimate_mass_thickness(
W_init = estimator.W_
H_init = estimator.H_

elts = list(self.model.get_elements(include_ignored=False))
elts_indices = self.model.NMF_simplex()
new_elts_dict = {elts[i]: W_init[elts_indices[i]] for i in range(len(elts))}
new_elts_dict = dict(zip(elts, W_init[elts_indices]))
total_weight = self._elements_dict_to_weights(new_elts_dict)
curr_mt = self._extract_mass_thickness(H_init.sum(), total_weight)

_ += 1

print(
f"The current estimated mass-thickness is {curr_mt} g.cm^-2",
flush=True,
Expand Down Expand Up @@ -677,7 +662,7 @@ def _elements_dict_to_weights(self, elements_dict):
Total weight of the elements in grams.
"""
total_weight = sum(
quantity * NPT["table"][element]["atomic_mass"] * 1.66053906660e-24
quantity * get_npt()[element]["atomic_mass"] * 1.66053906660e-24
for element, quantity in elements_dict.items()
)
return total_weight
Expand Down Expand Up @@ -736,7 +721,7 @@ def select_background_windows(self, num_windows=4, ranges=None):
cm = self._register_ranges
init_ranges = self._generate_ranges(num_windows)
self.spans = []
for i in range(num_windows):
for _ in range(num_windows):
self.spans.append(Signal1DRangeSelector(self))

for j, span in enumerate(self.spans):
Expand Down Expand Up @@ -766,7 +751,7 @@ def _register_ranges(self, signal, left, right):

def _compute_bremsstrahlung(self):
mt = self.metadata.Sample.density * self.metadata.Sample.thickness
elts_dict = {elt: 1.0 for elt in self.metadata.Sample.elements}
elts_dict = dict.fromkeys(self.metadata.Sample.elements, 1.0)
brstlg_model, mask = self.model.bremsstrahlung_only_tools(
mass_thickness=mt, elements_dict=elts_dict, ranges=self.ranges
)
Expand Down Expand Up @@ -798,10 +783,8 @@ def _plot_background(self, model):

def _generate_ranges(self, num):
axis = self.axes_manager.signal_axes[0]
bounds = (axis.low_value, axis.high_value)
values = np.linspace(bounds[0], bounds[1], num=2 * num + 2)
ranges_list = [(values[2 * i - 1], values[2 * i]) for i in range(1, num + 1)]
return ranges_list
values = np.linspace(axis.low_value, axis.high_value, num=2 * num + 2)
return values[1:-1].reshape(num, 2).tolist()

@perrymok perrymok Jul 30, 2026

Copy link
Copy Markdown
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

np.reshape can be a bit hard to understand


def decomposition(
self,
Expand Down Expand Up @@ -996,8 +979,7 @@ def concentration_report(self, selected_elts=[], W_input=None, fit_error=True):
G = self.learning_results.decomposition_algorithm.G_
H = self.learning_results.decomposition_algorithm.H_
N = get_explained_intensity_W(G, W, H)
sqN = np.sqrt(N)
percentages = sqN / N * 100
percentages = 100.0 / np.sqrt(N)

else:
W = W_input
Expand Down Expand Up @@ -1056,44 +1038,39 @@ def estimate_best_binning(self, inspect=False):
L = self.axes_manager[2].size
K = self.axes_manager[0].size * self.axes_manager[1].size

facx = np.arange(1, self.axes_manager[0].size // 2 + 1)
facy = np.arange(1, self.axes_manager[1].size // 2 + 1)
size_x = self.axes_manager[0].size
size_y = self.axes_manager[1].size

facx = np.arange(1, size_x // 2 + 1)
facy = np.arange(1, size_y // 2 + 1)

binx = [self.axes_manager[0].size / i for i in facx]
biny = [self.axes_manager[1].size / i for i in facy]
vars_est = np.array([])
biases_est = np.array([])
bprod = []
for i in zip(binx, biny):
bprod.append((i[0], i[1]))
binx = size_x / facx
biny = size_y / facy

for i in tqdm(bprod):
vars_est = []
biases_est = []

bprod = list(zip(binx, biny))
data = self.data
mean_d = np.mean(data)

for i0, i1 in tqdm(bprod):
# Bin the measurement dataset and upsample to bring it back to its original dimensionality
B = i[0] * i[1]
binned = self.rebin(scale=(i[0], i[1], 1))
upsampled = binned.rebin(
new_shape=(
self.axes_manager[0].size,
self.axes_manager[1].size,
self.axes_manager[2].size,
)
)
B = i0 * i1
binned = self.rebin(scale=(i0, i1, 1))
upsampled = binned.rebin(new_shape=(size_x, size_y, L))
upsampled_data = upsampled.data
data = self.data

# Estimator of variance (Lemma 4.3) - \widehat{Var} (\hat{y}_i) = \alpha ^2 y_{i}+ (1-\alpha)^2 \sum_{k \in \mathcal{K}} (w_k^2 n_{i,k})
vars_est = np.append(vars_est, np.mean(upsampled_data * 1 / B))
vars_est.append(np.mean(upsampled_data) / B)

# Estimator of squared bias (Lemma 4.4) - \widehat{Bias^2}(\hat{y}_i) = (1-\alpha)^2\left((y_{n_i} - y_{i})^2 - \sum_{k\in \mathcal{K}} w_k^2y_{i,k} - y_{i} \right)
biases_est = np.append(
biases_est,
np.mean(
(data - upsampled_data) ** 2
- 1 / B * upsampled_data
- (1 - 2 / B) * data
),
)
mean = np.mean((data - upsampled_data) ** 2 - upsampled_data / B)

biases_est.append(mean - (1 - 2 / B) * mean_d)

vars_est = np.array(vars_est)
biases_est = np.array(biases_est)
mprimes_est = vars_est * K / L + biases_est
estimated_binning = (
bprod[np.argmin(mprimes_est)][0],
Expand Down Expand Up @@ -1159,8 +1136,8 @@ def generate_part_fixed_H_matrix(
part_f_H : np.ndarray
A fixed H matrix for one phase.
"""
part_f_H = (-1) * np.ones(
shape=(self.data.shape[0], self.data.shape[1]), dtype=float
part_f_H = np.full(
shape=(self.data.shape[0], self.data.shape[1]), fill_value=-1.0, dtype=float
)

if value < 0:
Expand Down Expand Up @@ -1205,13 +1182,7 @@ def set_fixed_H(self, areas_dict: dict[str, np.ndarray]) -> np.ndarray:
A fixed H matrix for the SmoothNMF decomposition algorithm.
"""

H = (-1) * np.ones(
shape=(len(areas_dict), self.data.shape[0], self.data.shape[1]), dtype=float
)

for i, p in enumerate(areas_dict):
H[i, :, :] = areas_dict[p]

H = np.array(list(areas_dict.values()))

@perrymok perrymok Jul 30, 2026

Copy link
Copy Markdown
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

this is a pretty cool vectorisation.

return H.reshape((len(areas_dict), self.data.shape[0] * self.data.shape[1]))

@_check_decomposition
Expand Down Expand Up @@ -1247,20 +1218,16 @@ def num_elts(elements=[]):
_elts_indices = self.model.NMF_simplex()

skipped_elts = sym_elts(elements=skipped_elements)
elts = []
elts_indices = []
for i, elt in enumerate(_elts):
if elt in skipped_elts:
pass
else:
elts.append(elt)
elts_indices.append(_elts_indices[i])
elts = [elt for elt in _elts if elt not in skipped_elts]
elts_indices = [
idx for idx, elt in zip(_elts_indices, _elts) if elt not in skipped_elts
]

W = _W[elts_indices, :]
WH = W @ H
WH = WH.reshape([W.shape[0]] + list(self.data.shape[:-1]))

WH /= WH.sum(0)[np.newaxis, ...] / 100
WH /= WH.sum(axis=0, keepdims=True) / 100.0

if self.axes_manager.navigation_dimension == 2:
Signal = hs.signals.Signal2D
Expand All @@ -1282,7 +1249,8 @@ def num_elts(elements=[]):
self.axes_manager[i], ["units", "scale", "name", "offset"]
)
q.metadata.Signal.quantity = "Atomic %"
wh = Signal(WH)

wh = Signal(WH)

for i in range(self.axes_manager.navigation_dimension):
wh.axes_manager[1 + i].update_from(
Expand Down Expand Up @@ -1954,9 +1922,7 @@ def plot_table(
hover_data = []

cmap = plt.get_cmap("tab10")
# TODO: use cache when optimisation is merged
with open(SYMBOLS_PERIODIC_TABLE, "r") as f:
SPT = json.load(f)["table"]
SPT = get_spt()

for i, elt in enumerate(elements):
if (anum := str(SPT[elt]["number"])) not in table:
Expand Down
Loading