From dfb3150d07ac00e90a33c2ca39e0b19b2a98a023 Mon Sep 17 00:00:00 2001 From: Daniel Paley Date: Fri, 20 Mar 2026 20:24:07 -0400 Subject: [PATCH 01/15] small_cell: apply beamtime enhancements for subsampling and clustering Brings in the following enhancements from smx_beamtime: - New clustering tools (cluster2, index2) - Geometry refinement capabilities - Cake plot preparation utilities - Reflection subsampling and clustering in small_cell indexing - Various powder plotting enhancements These changes enable advanced small-cell crystallography workflows including geometry refinement and multi-stage clustering. --- .../small_cell/command_line/cake_plot_prep.py | 94 + xfel/small_cell/command_line/cluster2.py | 1182 ++++++++ xfel/small_cell/command_line/index2.py | 2640 +++++++++++++++++ .../command_line/powder_refine_geometry.py | 158 + xfel/small_cell/geometry_refiner.py | 321 ++ xfel/small_cell/powder_util.py | 182 +- xfel/small_cell/small_cell.py | 394 +-- 7 files changed, 4771 insertions(+), 200 deletions(-) create mode 100644 xfel/small_cell/command_line/cake_plot_prep.py create mode 100644 xfel/small_cell/command_line/cluster2.py create mode 100644 xfel/small_cell/command_line/index2.py create mode 100644 xfel/small_cell/command_line/powder_refine_geometry.py create mode 100644 xfel/small_cell/geometry_refiner.py diff --git a/xfel/small_cell/command_line/cake_plot_prep.py b/xfel/small_cell/command_line/cake_plot_prep.py new file mode 100644 index 00000000000..cdce3106704 --- /dev/null +++ b/xfel/small_cell/command_line/cake_plot_prep.py @@ -0,0 +1,94 @@ +from __future__ import division +# LIBTBX_SET_DISPATCHER_NAME cctbx.xfel.small_cell.cake_plot_prep +from dials.array_family import flex +from dxtbx.model.experiment_list import ExperimentList +from dials.array_family import flex +import sys, glob +import matplotlib.pyplot as plt + + +help_str = """ +Make a cake plot from DIALS spotfinder spots + +A cake plot is the azimuthal angle of a spot on an image vs. its resolution. +Powder rings will appear as vertical stripes, with defects in geometry +causing them to appear wavy. A cake plot is also insensitive to badly masked +regions of the detector compared to a 1d radial average as the aziumuthal +angle of a spot isn't averaged into the 1d trace. + +This script creates cake.npy which is used by +cctbx.xfel.small_cell.cake_plot. Run this script first to generate it, then +run cctbx.xfel.small_cell.cake_plot + +This script expects files named "*_strong.expt" and "_strong.refl". Supply +the former and the script will seek for the latter. + +Usage (note, wild cards are permitted, but quotes are recommended): +cctbx.xfel.small_cell.cake_plot_prep "/*_strong.expt>" + +Multiprocessing support is availible using MPI. Example: +mpirun cctbx.xfel.small_cell.cake_plot_prep "/*_strong.expt>" +""" + +def run(args): + if "-h" in args or "--help" in args: + if rank == 0: + print(help_str) + return + + filenames = [] + for arg in sys.argv[1:]: + filenames.extend(glob.glob(arg)) + if not filenames: + sys.exit("No data found") + + x, y = flex.double(), flex.double() + det = None + for fn in filenames: + print (fn) + #try: + refls = flex.reflection_table.from_file(fn.split('_strong.expt')[0] + "_strong.refl") + #except OSError: + # continue + expts = ExperimentList.from_file(fn, check_format=False) + for expt_id, expt in enumerate(expts): + subset = refls.select(expt_id == refls['id']) + if len(subset) > 200: continue + det = expt.detector + for panel_id, panel in enumerate(det): + r = subset.select(subset['panel'] == panel_id) + x_, y_, _ = r['xyzobs.px.value'].parts() + pix = panel.pixel_to_millimeter(flex.vec2_double(x_, y_)) + c = panel.get_lab_coord(pix) + x.extend(c.parts()[0]) + y.extend(c.parts()[1]) + + if det: + z = flex.double(len(x), sum([p.get_origin()[2] for p in det])/len(det)) + coords = flex.vec3_double(x,y,z) + two_theta = coords.angle((0,0,-1)) + d = expts[0].beam.get_wavelength() / 2 / flex.sin(two_theta/2) + azi = flex.vec3_double(x, y, flex.double(len(x), 0)).angle((0,1,0), deg=True) + azi.set_selected(x < 0, 180+(180-azi.select(x<0))) + else: + d = flex.double() + azi = flex.double() + + import numpy as np + fig, axes = plt.subplots(1, 1, figsize=(6, 3)) + axes.plot( + 1/d.as_numpy_array(), azi.as_numpy_array(), + marker='.', linestyle='none', + markersize=0.5, alpha=0.5 + ) + axes.set_ylabel('Azimuthal Angle') + axes.set_xlabel(r'Resolution ($\mathrm{\AA}$)') + x_ticks = np.array([10, 5, 2, 1]) + axes.set_xticks(1/x_ticks) + axes.set_xticklabels(x_ticks) + fig.tight_layout() + plt.show() + +if __name__ == "__main__": + run(sys.argv[1:]) + diff --git a/xfel/small_cell/command_line/cluster2.py b/xfel/small_cell/command_line/cluster2.py new file mode 100644 index 00000000000..4371312e1e6 --- /dev/null +++ b/xfel/small_cell/command_line/cluster2.py @@ -0,0 +1,1182 @@ +import numpy as np +import matplotlib.pyplot as plt +from matplotlib.widgets import SpanSelector, RectangleSelector +from matplotlib.patches import Rectangle +from scipy.stats import gaussian_kde +from scipy.ndimage import maximum_filter +from scipy.ndimage import generate_binary_structure +from sklearn.neighbors import KernelDensity +import sys +from scipy.optimize import minimize + +import matplotlib +matplotlib.use('TkAgg') + +def fit_gaussian_peak(points, initial_center, bandwidths): + """Fit 3D Gaussian to refine peak location""" + from scipy.optimize import minimize + + def negative_log_likelihood(params): + center = params[:3] + # Gaussian kernel with anisotropic bandwidths + diff = (points - center) / bandwidths + distances_sq = np.sum(diff**2, axis=1) + log_prob = -0.5 * distances_sq + return -np.sum(log_prob) # Negative for minimization + + result = minimize( + negative_log_likelihood, + initial_center, + method='BFGS' + ) + + if result.success: + return result.x + else: + return initial_center # Fall back to grid position + +def has_saddle_between(peak1, peak2, hist_smooth, edges, bin_width, prominence_threshold=0.5): + """Check if there's a significant saddle between two peaks""" + # Sample points along the line between peaks + n_samples = 50 + t = np.linspace(0, 1, n_samples) + line_points = peak1[np.newaxis, :] * (1 - t[:, np.newaxis]) + peak2[np.newaxis, :] * t[:, np.newaxis] + + # Convert to bin indices + bin_indices = [] + for dim in range(3): + indices = np.searchsorted(edges[dim], line_points[:, dim]) - 1 + indices = np.clip(indices, 0, hist_smooth.shape[dim] - 1) + bin_indices.append(indices) + + # Get density along the line + line_densities = hist_smooth[bin_indices[0], bin_indices[1], bin_indices[2]] + + # Find minimum along path + min_density = np.min(line_densities) + + # Get densities at the two peaks + peak1_idx = [np.searchsorted(edges[i], peak1[i]) - 1 for i in range(3)] + peak2_idx = [np.searchsorted(edges[i], peak2[i]) - 1 for i in range(3)] + peak1_idx = [np.clip(idx, 0, hist_smooth.shape[i] - 1) for i, idx in enumerate(peak1_idx)] + peak2_idx = [np.clip(idx, 0, hist_smooth.shape[i] - 1) for i, idx in enumerate(peak2_idx)] + + density1 = hist_smooth[tuple(peak1_idx)] + density2 = hist_smooth[tuple(peak2_idx)] + + # Check prominence: saddle should be significantly lower than both peaks + lower_peak = min(density1, density2) + prominence = (lower_peak - min_density) / lower_peak + + result = prominence > prominence_threshold + return result + + +class ManualClusterer: + def __init__(self, data, bandwidths=[0.001, 0.001, 1.0], n_maxima=500, + qvals_1=None, qvals_2=None, sb1_callback=None, recon=None, + qmin=.1, qmax=.5, points=None, n_shortest=500): + if len(data) > 1000000: + indices = np.random.choice(len(data), size=1000000, replace=False) + data = data[indices] + self.data = data + self.bandwidths = np.array(bandwidths) + self.n_maxima = n_maxima + self.n_shortest = n_shortest + self.qvals_1 = qvals_1 + self.qvals_2 = qvals_2 + self.sb1_qvals = None + self.qmin = qmin + self.qmax = qmax + self.points = points + + self.current_selection = None + self.final_selection = None + self.means = None + self.span_patch = None + self.rect_patch = None + self.selected_triplets = [] + + self.sb1_callback = sb1_callback + self.recon = recon + + # Compute KDE maxima + if n_maxima > 0: + self.compute_pf_maxima() + self.points = self.kde_maxima + + + def select_qvals(self, title="Select q-value ranges"): + """ + Simple q-value selection mode using 1D histogram on dimension 0. + + Parameters: + ----------- + title : str + Title for the histogram window + + Returns: + -------- + list : Selected q-values (medians of selected ranges) + """ + # Create figure for histogram + self.fig_qval, self.ax_qval = plt.subplots(num='Q-value Selection', figsize=(16, 4)) + + # Initialize list to store selected q-values + self.selected_qvals = [] + self.qval_spans = [] # Store span patches for visualization + + # Plot histogram + self._plot_qval_histogram(title) + + # Set up span selector + self.qval_span = SpanSelector( + self.ax_qval, + self._on_qval_span_select, + 'horizontal', + useblit=True, + props=dict(alpha=0.5, facecolor='blue') + ) + + # Add done button + done_ax = self.fig_qval.add_axes([0.9, 0.01, 0.09, 0.05]) + self.qval_done_button = plt.Button(done_ax, 'Done') + self.qval_done_button.on_clicked(self._on_qval_done) + + # Add clear button + clear_ax = self.fig_qval.add_axes([0.8, 0.01, 0.09, 0.05]) + self.qval_clear_button = plt.Button(clear_ax, 'Clear Last') + self.qval_clear_button.on_clicked(self._on_qval_clear_last) + + # Connect key press events + self.fig_qval.canvas.mpl_connect('key_press_event', self._on_qval_key_press) + + # Show plot and wait for user interaction + plt.show(block=True) + + # Close figure + plt.close(self.fig_qval) + + return self.selected_qvals + + def _plot_qval_histogram(self, title): + """Plot histogram for q-value selection.""" + self.ax_qval.clear() + + # Create histogram + plot_range = .1, self.qmax + hist, edges = np.histogram(self.data[:, 0], bins=2000, range=plotrange) + self.ax_qval.plot(edges[:-1], hist) + self.ax_qval.set_title('abcd') + self.ax_qval.set_xlabel('q-value') + self.ax_qval.set_ylabel('Count') + + # Add tick marks for previously selected q-values + if self.qvals_1 is not None: + self.ax_qval.plot(self.qvals_1, np.zeros_like(self.qvals_1), + '|', color='blue', markersize=25, markeredgewidth=2, label='qvals_1') + if self.qvals_2 is not None: + self.ax_qval.plot(self.qvals_2, np.zeros_like(self.qvals_2), + '|', color='red', markersize=25, markeredgewidth=2, label='qvals_2') + + # Plot already selected q-values from this session + if self.selected_qvals: + self.ax_qval.plot(self.selected_qvals, np.zeros_like(self.selected_qvals), + 'o', color='purple', markersize=10, label='Selected') + + # Redraw span patches + for patch in self.qval_spans: + self.ax_qval.add_patch(patch) + + if any([self.qvals_1 is not None, self.qvals_2 is not None, self.selected_qvals]): + self.ax_qval.legend() + + self.fig_qval.canvas.draw_idle() + + def _on_qval_span_select(self, xmin, xmax): + """Handle span selection for q-values.""" + # Select data within the span + mask = (self.data[:, 0] >= xmin) & (self.data[:, 0] <= xmax) + selected_data = self.data[mask, 0] + + if len(selected_data) > 0: + # Calculate median of selected range + median_qval = np.median(selected_data) + + # Add to selected q-values + self.selected_qvals.append(median_qval) + + # Create span patch for visualization + ylims = self.ax_qval.get_ylim() + span_patch = Rectangle((xmin, ylims[0]), xmax-xmin, ylims[1]-ylims[0], + alpha=0.2, color='purple') + self.qval_spans.append(span_patch) + + # Update plot + self._plot_qval_histogram(self.ax_qval.get_title()) + + print(f"Selected q-value: {median_qval:.6f} (from {len(selected_data)} points in range [{xmin:.6f}, {xmax:.6f}])") + + # Save to file + with open('selected_qvals.txt', 'a') as f: + f.write(f"{median_qval:.6f}\n") + + def _on_qval_clear_last(self, event): + """Remove the last selected q-value.""" + if self.selected_qvals: + removed = self.selected_qvals.pop() + if self.qval_spans: + self.qval_spans.pop() + print(f"Removed q-value: {removed:.6f}") + + # Update plot + self._plot_qval_histogram(self.ax_qval.get_title()) + + def _on_qval_key_press(self, event): + """Handle key press events for q-value selection.""" + if event.key == 'c': + # Clear last selection (same as button) + self._on_qval_clear_last(None) + elif event.key == 'd' or event.key == 'enter': + # Done selecting (same as button) + self._on_qval_done(None) + elif event.key == 'escape': + # Cancel and close without saving + self.selected_qvals = [] + plt.close(self.fig_qval) + + def _on_qval_done(self, event): + """Called when Done button is clicked for q-value selection.""" + print(f"Selected {len(self.selected_qvals)} q-values: {self.selected_qvals}") + plt.close(self.fig_qval) + + def select_triplets(self, title=None): + """Run the interactive selection process and return the selected triplets.""" + # Create the three windows with specific sizes + self.fig1, self.ax1 = plt.subplots(num='Step 1: Histogram Selection', figsize=(16,3)) + self.fig2, self.ax2 = plt.subplots(num='Step 2: 2D Selection', figsize=(16,3)) + self.fig3, (self.ax3a, self.ax3b, self.ax3c) = plt.subplots(1, 3, num='Step 3: Final Selection', + figsize=(16,3)) + + # Set window positions to stack them vertically + backend = plt.get_backend() + assert 'Tk' in backend + manager1 = self.fig1.canvas.manager + manager2 = self.fig2.canvas.manager + manager3 = self.fig3.canvas.manager + dpi = self.fig1.dpi + height1 = int(3 * dpi) # 3 inches * dpi + + # Position windows with some spacing + self.fig1.canvas.manager.window.wm_geometry("+100+50") + self.fig2.canvas.manager.window.wm_geometry(f"+100+{50 + height1 + 40}") + self.fig3.canvas.manager.window.wm_geometry(f"+100+{50 + 2*height1 + 80}") + + # Set fixed subplot sizes + #self.fig3.set_tight_layout(False) + self.fig1.subplots_adjust(bottom=.2) + self.fig2.subplots_adjust(bottom=.2) + self.fig3.subplots_adjust(bottom=.2) + + # Initialize the first window + self.show_histogram(title=title) + + # Set up the done button + done_ax = self.fig3.add_axes([0.9, 0.01, 0.09, 0.05]) + self.done_button = plt.Button(done_ax, 'Done') + self.done_button.on_clicked(self.on_done) + + # Flag to track when selection is complete + self.selection_done = False + + # Show plots and wait for user interaction + plt.show(block=True) + + # Close all figures + plt.close(self.fig1) + plt.close(self.fig2) + plt.close(self.fig3) + + # Return the selected triplets + return self.selected_triplets + + def compute_kde_maxima(self): + # Normalize data by bandwidths for anisotropic KDE + normalized_data = self.data / self.bandwidths[np.newaxis, :] + + # Create KernelDensity object + kde = KernelDensity(bandwidth=1, kernel='cosine') + print('fit') + kde.fit(normalized_data) + print('done fit') + + # Take a random subsample (5%) for evaluation + n_sample = max(int(len(normalized_data) * 0.05), 100000) # At least 1000 points + n_sample = 200000 + n_sample = min(n_sample, len(normalized_data)) # Can't sample more than we have + print(f'{n_sample=}') + + # Random sampling without replacement + sample_indices = np.random.choice(len(normalized_data), size=n_sample, replace=False) + sample_data = normalized_data[sample_indices] + + # Evaluate KDE at sampled points + print('score') + sample_densities = np.exp(kde.score_samples(sample_data)) + print('done score') + + # Sort sampled points by density value + sorted_indices = np.argsort(sample_densities)[::-1] + + # Filter to avoid maxima that are too close together + safety_factor = 2.0 # How many bandwidths apart maxima should be + min_distances = self.bandwidths * safety_factor + + # Initialize list to store maxima + maxima_indices = [] + maxima_points = [] + maxima_densities = [] + + # Function to be minimized (negative density) + def negative_density(x): + return -np.exp(kde.score_samples([x])[0]) + + # Process points in order of decreasing density + for i, idx in enumerate(sorted_indices): + if i%1000==0: + print('processed', i, ', kept', len(maxima_points)) + if len(maxima_points) >= self.n_maxima: + break + + initial_point = sample_data[idx] + q1, q2, th = initial_point + if np.abs(q1-q2) < 2: continue + if th<8 or th>160: continue + + optimized_point = initial_point + optimized_density = 1 +# # Run optimization to find true maximum +# result = minimize( +# negative_density, +# initial_point, +# method='BFGS', +# #options={'gtol': 1e-5} # Gradient tolerance for convergence +# ) +# +# if result.success: +# optimized_point = result.x +# optimized_density = -result.fun # Negate back to get positive density + + + # Check if this point is far enough from all accepted maxima + too_close = False + for accepted_point in maxima_points: + dist = np.linalg.norm(optimized_point - accepted_point) + if dist < 5: + too_close = True + break + + sus = False + q1,q2,th = optimized_point + if np.abs(q1-q2) < 2: sus = True + if th < 5 or th > 160: sus = True + + if not too_close and not sus: + maxima_points.append(optimized_point) + maxima_densities.append(optimized_density) + + # Store the maxima and their density values + self.kde_maxima = np.array(maxima_points) * self.bandwidths[np.newaxis, :] + self.kde_values = np.array(maxima_densities) + with open('autopeaks.txt', 'w') as f: + for x in self.kde_maxima: + print(round(x[0], 4), round(x[1], 4), round(x[2], 2), file=f) + + + print(f"Found {len(self.kde_maxima)} KDE maxima from {n_sample} sampled points") + + def compute_pf_maxima(self): + # Define bins with 2x oversampling + bin_width = self.bandwidths / 2.0 + + bins = [ + np.arange(0.05, 0.5 + bin_width[0], bin_width[0]), + np.arange(0.05, 0.5 + bin_width[1], bin_width[1]), + np.arange(10, 160 + bin_width[2], bin_width[2]) + ] + + print(f'Creating histogram with shape: {[len(b)-1 for b in bins]}') + + from scipy.spatial import cKDTree + + # Build tree once at start + tree = cKDTree(self.data / self.bandwidths) + + # Create 3D histogram + hist, edges = np.histogramdd(self.data, bins=bins) + + # Smooth with gaussian (sigma ~ 1 bin to merge nearby peaks) + from scipy import ndimage + hist_smooth = ndimage.gaussian_filter(hist, sigma=1) + + # Find local maxima + max_filtered = ndimage.maximum_filter(hist_smooth, size=10) + peaks = (hist_smooth == max_filtered) + + # Threshold to remove noise peaks + threshold = hist_smooth.max() * 0.01 # Adjust as needed + peaks &= (hist_smooth > threshold) + + print(f'Found {peaks.sum()} initial peaks') + + # Get peak coordinates in bin indices + peak_indices = np.argwhere(peaks) + peak_values = hist_smooth[peaks] + + # Convert bin indices to actual coordinates + peak_coords = np.array([ + edges[0][peak_indices[:, 0]] + bin_width[0]/2, + edges[1][peak_indices[:, 1]] + bin_width[1]/2, + edges[2][peak_indices[:, 2]] + bin_width[2]/2 + ]).T + + # Apply your domain filters + valid_mask = np.ones(len(peak_coords), dtype=bool) + + q1, q2, th = peak_coords[:, 0], peak_coords[:, 1], peak_coords[:, 2] + valid_mask &= (np.abs(q1 - q2) >= 2 * self.bandwidths[0]) + valid_mask &= (th >= 8) & (th <= 160) + + peak_coords = peak_coords[valid_mask] + peak_values = peak_values[valid_mask] + sortkey = peak_values + #sortkey = -1* (peak_coords[:,0] + peak_coords[:,1]) + + # Sort by density + sorted_indices = np.argsort(sortkey)[::-1] + + # Filter by minimum distance (your safety_factor logic) + min_dist = 10.0 # In your normalized space this was 5 + final_peaks = [] + final_values = [] + + for idx in sorted_indices: + if len(final_peaks) >= self.n_maxima: + break + + candidate = peak_coords[idx] + + # Check if it's a true local maximum + if len(final_peaks) > 0: + # Option 1: Saddle test + is_separate = all( + has_saddle_between( + candidate, accepted, hist_smooth, edges, bin_width, prominence_threshold=0.8 + ) + for accepted in final_peaks + if np.linalg.norm((candidate - accepted) / bin_width) < 180 # Normalized distance + ) + + # Option 2: Gradient test (faster) + # is_separate = is_local_maximum(candidate, hist_smooth, edges, bin_width, final_peaks) + + if not is_separate: + continue + + final_peaks.append(candidate) + final_values.append(peak_values[idx]) + + # Refine peak locations by fitting Gaussians + refined_peaks = [] + + for i, peak in enumerate(final_peaks): + if i % 100 == 0: + print(f'Refining peak {i}/{len(final_peaks)}') + + # Select points within 2x bandwidth + normalized_peak = peak / self.bandwidths + indices = tree.query_ball_point(normalized_peak, r=2.0) + nearby_points = self.data[indices] + + if len(nearby_points) >= 5: # Need enough points to fit + refined_peak = fit_gaussian_peak( + nearby_points, peak, self.bandwidths + ) + refined_peaks.append(refined_peak) + else: + print('fallback') + refined_peaks.append(peak) # Not enough points, keep grid location + + self.kde_maxima = np.array(refined_peaks) + self.kde_values = np.array(final_values) + + # Final sorting + sort_vals = self.kde_maxima[:, 0] + self.kde_maxima[:, 1] + final_indices = np.argsort(sort_vals)[:self.n_shortest] + self.kde_maxima = self.kde_maxima[final_indices] + self.kde_values = self.kde_values[final_indices] + + + def compute_slice_maxima(self, n_q1_peaks=30, min_q1=0.10, max_q1=0.35, + slice_width=0.002, max_peaks_per_slice=100, + max_peaks=500, verbose=True): + """ + Find cluster centers using slice-based 2D KDE approach. + + This method finds q1 peaks using 1D KDE, then performs 2D KDE + peak finding within each q1 slice, followed by 3D refinement. + + Final peaks are selected by score / sqrt(q1 * q2) to prioritize + low-q peaks which are most valuable for unit cell determination. + + Parameters + ---------- + n_q1_peaks : int + Maximum number of q1 slices to process + min_q1, max_q1 : float + Range for q1 peak detection + slice_width : float + Width of each q1 slice + max_peaks_per_slice : int + Maximum peaks to find per slice + verbose : bool + Print progress + """ + from scipy.ndimage import maximum_filter + from scipy.optimize import minimize_scalar + from scipy.spatial import cKDTree + + data = self.data + bandwidths_3d = self.bandwidths + bandwidths_2d = np.array([bandwidths_3d[1], bandwidths_3d[2]]) + + # Step 1: Find q1 peaks using 1D Gaussian KDE + if verbose: + print("Finding q1 peaks...") + + q1_data = data[:, 0] + q1_mask = (q1_data >= min_q1) & (q1_data <= max_q1) + q1_filtered = q1_data[q1_mask] + + # Subsample for KDE + np.random.seed(42) + n_sample = min(1000000, len(q1_filtered)) + sample_idx = np.random.choice(len(q1_filtered), n_sample, replace=False) + q1_sample = q1_filtered[sample_idx].reshape(-1, 1) + + # Fit 1D KDE + kde_1d = KernelDensity(bandwidth=0.0002, kernel='cosine') + kde_1d.fit(q1_sample) + + # Find peaks on fine grid + n_eval = 10000 + q1_grid = np.linspace(min_q1, max_q1, n_eval) + density_1d = np.exp(kde_1d.score_samples(q1_grid.reshape(-1, 1))) + + footprint = np.ones(5) + local_max = maximum_filter(density_1d, footprint=footprint) + peaks_mask = (density_1d == local_max) & (density_1d > np.percentile(density_1d, 50)) + peak_q1s_raw = q1_grid[np.where(peaks_mask)[0]] + + # Refine peak positions + def neg_density(q1): + return -kde_1d.score_samples([[q1]])[0] + + refined_q1s = [] + for q1 in peak_q1s_raw: + result = minimize_scalar(neg_density, bounds=(q1-0.0005, q1+0.0005), method='bounded') + refined_q1s.append(result.x) + refined_q1s = np.array(refined_q1s) + + # Score and select with spacing + peak_densities = np.exp(kde_1d.score_samples(refined_q1s.reshape(-1, 1))) + scores = peak_densities / (refined_q1s ** 0.5) + sort_idx = np.argsort(scores)[::-1] + q1s_sorted = refined_q1s[sort_idx] + + min_spacing = 0.002 + q1_centers = [] + for q1 in q1s_sorted: + if len(q1_centers) >= n_q1_peaks: + break + if all(abs(q1 - s) >= min_spacing for s in q1_centers): + q1_centers.append(q1) + q1_centers = np.sort(q1_centers) + + if verbose: + print(f"Found {len(q1_centers)} q1 peaks") + + # Step 2: Process each q1 slice + all_peaks = [] + all_scores = [] + + for i, q1_center in enumerate(q1_centers): + q1_min = q1_center - slice_width / 2 + q1_max = q1_center + slice_width / 2 + + slice_mask = (data[:, 0] >= q1_min) & (data[:, 0] <= q1_max) + slice_data = data[slice_mask] + + if len(slice_data) < 50: + continue + + if verbose: + print(f"Slice {i+1}/{len(q1_centers)}: q1=[{q1_min:.4f}, {q1_max:.4f}], {len(slice_data)} points", end="") + + # Filter diagonal and theta + diag_threshold = 2.0 + theta_range = (8, 160) + q_diff_norm = np.abs(slice_data[:, 0] - slice_data[:, 1]) / bandwidths_3d[0] + off_diag_mask = q_diff_norm >= diag_threshold + theta_mask = (slice_data[:, 2] >= theta_range[0]) & (slice_data[:, 2] <= theta_range[1]) + slice_data_filtered = slice_data[off_diag_mask & theta_mask] + + if len(slice_data_filtered) < 50: + if verbose: + print(" -> 0 peaks (filtered)") + continue + + # 3D KDE and tree on full slice + normalized_3d = slice_data / bandwidths_3d + kde_3d = KernelDensity(bandwidth=1, kernel='cosine') + kde_3d.fit(normalized_3d) + tree_3d = cKDTree(normalized_3d) + + # 2D KDE on filtered subsample + max_sample_2d = 30000 + if len(slice_data_filtered) > max_sample_2d: + sample_idx_2d = np.random.choice(len(slice_data_filtered), max_sample_2d, replace=False) + slice_2d_sample = slice_data_filtered[sample_idx_2d, 1:3] + else: + slice_2d_sample = slice_data_filtered[:, 1:3] + + normalized_2d_sample = slice_2d_sample / bandwidths_2d + kde_2d = KernelDensity(bandwidth=1, kernel='cosine') + kde_2d.fit(normalized_2d_sample) + scores_2d = kde_2d.score_samples(normalized_2d_sample) + tree_2d = cKDTree(normalized_2d_sample) + + # Find 2D peaks + def refine_2d(point_norm): + current = point_norm.copy() + for _ in range(5): + idx = tree_2d.query_ball_point(current, r=1.5) + if len(idx) < 5: + break + nearby = normalized_2d_sample[idx] + nearby_scores = scores_2d[idx] + weights = np.exp(nearby_scores - nearby_scores.max()) + new_pos = np.average(nearby, axis=0, weights=weights) + if np.linalg.norm(new_pos - current) < 0.01: + break + current = new_pos + return current + + top_k = min(5000, len(slice_2d_sample)) + top_idx = np.argsort(scores_2d)[-top_k:] + + min_distance_2d = 3.0 + peaks_2d_norm = [] + for idx in np.argsort(scores_2d[top_idx])[::-1]: + point_idx = top_idx[idx] + refined = refine_2d(normalized_2d_sample[point_idx]) + is_close = any(np.linalg.norm(refined - ex) < min_distance_2d for ex in peaks_2d_norm) + if not is_close: + peaks_2d_norm.append(refined) + if len(peaks_2d_norm) >= max_peaks_per_slice: + break + + # Refine in 3D + def refine_3d(point_3d, max_neighbors=300): + current = point_3d.copy() / bandwidths_3d + for _ in range(10): + idx = tree_3d.query_ball_point(current, r=2.0) + if len(idx) < 5: + break + idx = np.array(idx) + if len(idx) > max_neighbors: + idx = idx[np.random.choice(len(idx), max_neighbors, replace=False)] + nearby = normalized_3d[idx] + nearby_scores = kde_3d.score_samples(nearby) + weights = np.exp(nearby_scores - nearby_scores.max()) + new_pos = np.average(nearby, axis=0, weights=weights) + if np.linalg.norm(new_pos - current) < 0.01: + break + current = new_pos + return current * bandwidths_3d + + min_distance_3d = 3.0 + slice_peaks = [] + slice_scores = [] + + for p2d_norm in peaks_2d_norm: + p2d = p2d_norm * bandwidths_2d + dq2 = np.abs(slice_data[:, 1] - p2d[0]) + dtheta = np.abs(slice_data[:, 2] - p2d[1]) + nearby_mask = (dq2 < 0.003) & (dtheta < 4.0) + + if nearby_mask.sum() < 5: + continue + + nearby_scores_3d = kde_3d.score_samples(normalized_3d[nearby_mask]) + best_nearby = slice_data[nearby_mask][np.argmax(nearby_scores_3d)] + + refined_3d = refine_3d(best_nearby) + refined_norm = refined_3d / bandwidths_3d + + # Check theta range + if refined_3d[2] < theta_range[0] or refined_3d[2] > theta_range[1]: + continue + + # Check distance to existing peaks + is_close = any(np.linalg.norm(refined_norm - (ex / bandwidths_3d)) < min_distance_3d + for ex in slice_peaks) + if is_close: + continue + + refined_score = kde_3d.score_samples(refined_norm.reshape(1, -1))[0] + slice_peaks.append(refined_3d) + slice_scores.append(refined_score) + + all_peaks.extend(slice_peaks) + all_scores.extend(slice_scores) + + if verbose: + print(f" -> {len(slice_peaks)} peaks") + + # Deduplicate across slices + if len(all_peaks) > 0: + all_peaks = np.array(all_peaks) + all_scores = np.array(all_scores) + + sort_idx = np.argsort(all_scores)[::-1] + all_peaks = all_peaks[sort_idx] + all_scores = all_scores[sort_idx] + + unique_peaks = [] + unique_scores = [] + q_tol, theta_tol = 0.0005, 2.0 + + for p, s in zip(all_peaks, all_scores): + is_dup = any( + abs(p[0] - u[0]) < q_tol and + abs(p[1] - u[1]) < q_tol and + abs(p[2] - u[2]) < theta_tol + for u in unique_peaks + ) + if not is_dup: + unique_peaks.append(p) + unique_scores.append(s) + + unique_peaks = np.array(unique_peaks) + unique_scores = np.array(unique_scores) + + # Select top peaks prioritizing low q1*q2 (log-space adjustment) + if max_peaks is not None and len(unique_peaks) > max_peaks: + q1_vals = unique_peaks[:, 0] + q2_vals = unique_peaks[:, 1] + # score - 0.5*log(q1*q2) is the log-space equivalent of score/sqrt(q1*q2) + selection_scores = unique_scores - 0.5 * np.log(q1_vals * q2_vals) + top_idx = np.argsort(selection_scores)[-max_peaks:] + unique_peaks = unique_peaks[top_idx] + unique_scores = unique_scores[top_idx] + + if verbose: + print(f"Selected top {max_peaks} peaks (prioritizing low-q)") + + self.kde_maxima = unique_peaks + self.kde_values = unique_scores + else: + self.kde_maxima = np.array([]).reshape(0, 3) + self.kde_values = np.array([]) + + if verbose: + print(f"Final: {len(self.kde_maxima)} maxima") + + def show_histogram(self, title=None): + self.ax1.clear() + plotrange=self.qmin, self.qmax + self.ax1.set_xlim(plotrange) + hist, edges = np.histogram(self.data[:,0], bins=2000, range=plotrange) + self.ax1.plot(edges[:-1], hist) + if title is None: + title = "Select range in histogram" + self.ax1.set_title(title) + self.ax1.set_xlabel('q1') + self.ax1.set_ylabel('counts') + + # Add tick marks for KDE maxima if available + if hasattr(self, 'kde_maxima'): + self.ax1.plot(self.kde_maxima[:, 0], np.zeros_like(self.kde_maxima[:, 0]), + '|', color='green', markersize=20) + + # Add tick marks for specified q-values + if self.qvals_1 is not None: + self.ax1.plot(self.qvals_1, np.zeros_like(self.qvals_1), + '|', color='blue', markersize=25, markeredgewidth=2) + if self.qvals_2 is not None: + self.ax1.plot(self.qvals_2, np.zeros_like(self.qvals_2), + '|', color='red', markersize=25, markeredgewidth=2) + + if self.sb1_qvals is not None: + self.ax1.plot(self.sb1_qvals, np.zeros_like(self.sb1_qvals), + '|', color='orange', markersize=25, markeredgewidth=2) + + + self.span = SpanSelector( + self.ax1, + self.on_span_select, + 'horizontal', + useblit=True, + props=dict(alpha=0.5, facecolor='red') + ) + + # Redraw the figure + self.fig1.canvas.draw_idle() + + def on_span_select(self, xmin, xmax): + # Remove previous span patch if it exists + if self.span_patch is not None: + try: + self.span_patch.remove() + except NotImplementedError: + pass + self.span_patch = None + + # Create new span patch + ylims = self.ax1.get_ylim() + self.span_patch = Rectangle((xmin, ylims[0]), xmax-xmin, ylims[1]-ylims[0], + alpha=0.2, color='red') + self.ax1.add_patch(self.span_patch) + + # Select data within the span + mask = (self.data[:, 0] >= xmin) & (self.data[:, 0] <= xmax) + self.current_selection = self.data[mask] + + # Select KDE maxima within the span + if hasattr(self, 'kde_maxima'): + kde_mask = (self.kde_maxima[:, 0] >= xmin) & (self.kde_maxima[:, 0] <= xmax) + self.current_kde_selection = self.kde_maxima[kde_mask] + + # Select predicted points in the span + if self.points is not None: + points_mask = (self.points[:,0] >= xmin) & (self.points[:,0] <= xmax) + self.current_points_selection = self.points[points_mask] + else: + self.current_points_selection = None + + # Show the second window + self.show_scatter_2d() + + # Redraw both figures + self.fig1.canvas.draw_idle() + self.fig2.canvas.draw_idle() + + def show_scatter_2d(self): + self.ax2.clear() + plotrange = self.qmin, self.qmax + self.ax2.set_xlim(plotrange) + self.ax2.set_xlabel('q2') + self.ax2.set_ylabel('theta') + if self.current_selection is not None and len(self.current_selection) > 0: + self.ax2.scatter( + self.current_selection[:, 1], + self.current_selection[:, 2], + s=2, alpha=.2) + self.ax2.set_title(f'Draw box to select points (n={len(self.current_selection)})') + + # Plot KDE maxima if available + if hasattr(self, 'current_kde_selection') and len(self.current_kde_selection) > 0: + self.ax2.scatter(self.current_kde_selection[:, 1], self.current_kde_selection[:, 2], + color='red', marker='x', s=50, label='KDE maxima') + + + # Add vertical lines for q-values (second dimension) + if self.qvals_2 is not None: + ylim = self.ax2.get_ylim() + for q in self.qvals_2: + if q >= self.current_selection[:, 1].min() and q <= self.current_selection[:, 1].max(): + self.ax2.axvline(q, color='blue', linestyle='--', alpha=0.5) + if self.qvals_1 is not None and self.points is None: + for q in self.qvals_1: + if q >= self.current_selection[:, 1].min() and q <= self.current_selection[:, 1].max(): + self.ax2.axvline(q, color='red', linestyle='--', alpha=0.5) + if self.current_points_selection is not None: + self.ax2.scatter(self.current_points_selection[:,1], self.current_points_selection[:,2], + color='orange', marker='o', s=20, label='points') + + # Create RectangleSelector only if it doesn't exist already + if not hasattr(self, 'rect') or self.rect is None: + self.rect = RectangleSelector( + self.ax2, + self.on_rect_select, + useblit=True, + props=dict(facecolor='red', alpha=0.2) + ) + + # If there was a previous rectangle, redraw it + if self.rect_patch is not None: + self.ax2.add_patch(self.rect_patch) + else: + self.ax2.set_title('No points selected') + + self.fig2.canvas.mpl_connect('key_press_event', self.on_key_press) + self.fig2.canvas.draw_idle() + + def on_rect_select(self, eclick, erelease): + x1, y1 = eclick.xdata, eclick.ydata + x2, y2 = erelease.xdata, erelease.ydata + + # Remove previous rectangle if it exists + if self.rect_patch is not None: + self.rect_patch.remove() + + # Create new rectangle patch + self.rect_patch = Rectangle((min(x1, x2), min(y1, y2)), + abs(x2-x1), abs(y2-y1), + alpha=0.2, color='red') + self.ax2.add_patch(self.rect_patch) + self.fig2.canvas.draw_idle() + + # Select data within the rectangle + mask = ( + (self.current_selection[:, 1] >= min(x1, x2)) & + (self.current_selection[:, 1] <= max(x1, x2)) & + (self.current_selection[:, 2] >= min(y1, y2)) & + (self.current_selection[:, 2] <= max(y1, y2)) + ) + self.final_selection = self.current_selection[mask] + + # Select KDE maxima within the rectangle + if hasattr(self, 'current_kde_selection'): + kde_mask = ( + (self.current_kde_selection[:, 1] >= min(x1, x2)) & + (self.current_kde_selection[:, 1] <= max(x1, x2)) & + (self.current_kde_selection[:, 2] >= min(y1, y2)) & + (self.current_kde_selection[:, 2] <= max(y1, y2)) + ) + self.final_kde_selection = self.current_kde_selection[kde_mask] + + # Show the third window + self.show_final_plots() + + def show_final_plots(self): + self.update_means() + + # Clear all axes + for ax in [self.ax3a, self.ax3b, self.ax3c]: + ax.clear() + #ax.set_aspect('equal', adjustable='datalim') + + # Find the limits that contain all selected points + x_min, x_max = self.final_selection[:, 0].min(), self.final_selection[:, 0].max() + y_min, y_max = self.final_selection[:, 1].min(), self.final_selection[:, 1].max() + z_min, z_max = self.final_selection[:, 2].min(), self.final_selection[:, 2].max() + + # Add a small margin + margin = 0.1 + x_range = x_max - x_min + y_range = y_max - y_min + z_range = z_max - z_min + +# x_min -= margin * x_range +# x_max += margin * x_range +# y_min -= margin * y_range +# y_max += margin * y_range +# z_min -= margin * z_range +# z_max += margin * z_range + + # Set the limits for each plot + self.ax3a.set_xlim(x_min, x_max) + self.ax3a.set_ylim(y_min, y_max) + self.ax3a.set_xlabel('q1') + self.ax3a.set_ylabel('q2') + self.ax3b.set_xlabel('q1') + self.ax3b.set_ylabel('theta') + self.ax3c.set_xlabel('q2') + self.ax3c.set_ylabel('theta') + + self.ax3b.set_xlim(x_min, x_max) + self.ax3b.set_ylim(z_min, z_max) + + self.ax3c.set_xlim(y_min, y_max) + self.ax3c.set_ylim(z_min, z_max) + + # Plot all points within the final selection + self.ax3a.scatter(self.final_selection[:, 0], self.final_selection[:, 1], color='tab:blue', + s=5, alpha=.3) + self.ax3b.scatter(self.final_selection[:, 0], self.final_selection[:, 2], color='tab:blue', + s=5, alpha=.3) + self.ax3c.scatter(self.final_selection[:, 1], self.final_selection[:, 2], color='tab:blue', + s=5, alpha=.3) + + # Plot KDE maxima + if hasattr(self, 'final_kde_selection') and len(self.final_kde_selection) > 0: + self.ax3a.scatter(self.final_kde_selection[:, 0], self.final_kde_selection[:, 1], + color='red', marker='x', s=50) + self.ax3b.scatter(self.final_kde_selection[:, 0], self.final_kde_selection[:, 2], + color='red', marker='x', s=50) + self.ax3c.scatter(self.final_kde_selection[:, 1], self.final_kde_selection[:, 2], + color='red', marker='x', s=50) + + # Plot means + self.ax3a.plot(self.means[0], self.means[1], 'ro', markersize=6) + self.ax3b.plot(self.means[0], self.means[2], 'ro', markersize=6) + self.ax3c.plot(self.means[1], self.means[2], 'ro', markersize=6) + + # Add rectangle selectors to all plots + self.rect_final = [ + RectangleSelector(ax, self.on_final_rect_select, useblit=True, + props=dict(facecolor='red', alpha=0.2)) + for ax in [self.ax3a, self.ax3b, self.ax3c] + ] + + self.fig3.canvas.mpl_connect('key_press_event', self.on_key_press) + self.fig3.canvas.draw_idle() + + + + def on_final_rect_select(self, eclick, erelease): + x1, y1 = eclick.xdata, eclick.ydata + x2, y2 = erelease.xdata, erelease.ydata + + # Get the current axis + ax = eclick.inaxes + + # Create mask based on which plot was clicked + if ax == self.ax3a: + mask = ( + (self.final_selection[:, 0] >= min(x1, x2)) & + (self.final_selection[:, 0] <= max(x1, x2)) & + (self.final_selection[:, 1] >= min(y1, y2)) & + (self.final_selection[:, 1] <= max(y1, y2)) + ) + elif ax == self.ax3b: + mask = ( + (self.final_selection[:, 0] >= min(x1, x2)) & + (self.final_selection[:, 0] <= max(x1, x2)) & + (self.final_selection[:, 2] >= min(y1, y2)) & + (self.final_selection[:, 2] <= max(y1, y2)) + ) + elif ax == self.ax3c: + mask = ( + (self.final_selection[:, 1] >= min(x1, x2)) & + (self.final_selection[:, 1] <= max(x1, x2)) & + (self.final_selection[:, 2] >= min(y1, y2)) & + (self.final_selection[:, 2] <= max(y1, y2)) + ) + + # Save the limits + axl = self.ax3a.get_xlim() + ayl = self.ax3a.get_ylim() + bxl = self.ax3b.get_xlim() + byl = self.ax3b.get_ylim() + cxl = self.ax3c.get_xlim() + cyl = self.ax3c.get_ylim() + # Clear all plots including mean markers + for ax in [self.ax3a, self.ax3b, self.ax3c]: + ax.clear() + #ax.set_aspect('equal', adjustable='datalim') + # Restore the original limits + if ax == self.ax3a: + ax.set_xlim(axl) + ax.set_ylim(ayl) + elif ax == self.ax3b: + ax.set_xlim(bxl) + ax.set_ylim(byl) + else: + ax.set_xlim(cxl) + ax.set_ylim(cyl) + + # Update means based on the new selection + selected_points = self.final_selection[mask] + self.means = np.mean(selected_points, axis=0) + + # Replot everything + self.ax3a.scatter(self.final_selection[~mask, 0], self.final_selection[~mask, 1], + color='gray', s=5, alpha=0.3) + self.ax3a.scatter(self.final_selection[mask, 0], self.final_selection[mask, 1], + s=5, alpha=.3, color='tab:blue') + self.ax3a.plot(self.means[0], self.means[1], 'ro', markersize=6) + + self.ax3b.scatter(self.final_selection[~mask, 0], self.final_selection[~mask, 2], + color='gray', s=5, alpha=0.3) + self.ax3b.scatter(self.final_selection[mask, 0], self.final_selection[mask, 2], + s=5, alpha=.3, color='tab:blue') + self.ax3b.plot(self.means[0], self.means[2], 'ro', markersize=6) + + self.ax3c.scatter(self.final_selection[~mask, 1], self.final_selection[~mask, 2], + color='gray', s=5, alpha=0.3) + self.ax3c.scatter(self.final_selection[mask, 1], self.final_selection[mask, 2], + s=5, alpha=.3, color='tab:blue') + self.ax3c.plot(self.means[1], self.means[2], 'ro', markersize=6) + + # Plot KDE maxima + if hasattr(self, 'final_kde_selection') and len(self.final_kde_selection) > 0: + self.ax3a.scatter(self.final_kde_selection[:, 0], self.final_kde_selection[:, 1], + color='red', marker='x', s=50) + self.ax3b.scatter(self.final_kde_selection[:, 0], self.final_kde_selection[:, 2], + color='red', marker='x', s=50) + self.ax3c.scatter(self.final_kde_selection[:, 1], self.final_kde_selection[:, 2], + color='red', marker='x', s=50) + + # Plot means + self.ax3a.plot(self.means[0], self.means[1], 'ro', markersize=6) + self.ax3b.plot(self.means[0], self.means[2], 'ro', markersize=6) + self.ax3c.plot(self.means[1], self.means[2], 'ro', markersize=6) + + self.ax3a.set_xlabel('q1') + self.ax3a.set_ylabel('q2') + self.ax3b.set_xlabel('q1') + self.ax3b.set_ylabel('theta') + self.ax3c.set_xlabel('q2') + self.ax3c.set_ylabel('theta') + + self.fig3.canvas.draw_idle() + + def update_means(self): + self.means = np.mean(self.final_selection, axis=0) + + def on_key_press(self, event): + if event.key == 'a': + # Add mean to selected triplets + triplet = [self.means[0], self.means[1], self.means[2]] + self.selected_triplets.append(triplet) + + print(f"Selected point {self.means[0]:.6f} {self.means[1]:.6f} {self.means[2]:.6f}. " + f"Total triplets: {len(self.selected_triplets)}") + + # Also append to file if desired + with open('cluster_means.txt', 'a') as f: + np.savetxt(f, [self.means], fmt='%.6f') + + if self.sb1_callback is not None: + self.sb1_qvals, self.points = self.sb1_callback(triplet, self.recon) + self.show_histogram(title=self.ax1.get_title()) + + def on_done(self, event): + """Called when the Done button is clicked.""" + self.selection_done = True + plt.close('all') # Close all open figures + import gc;gc.collect() + + +# Example usage: +if __name__ == "__main__": + # Janky format + alldata = [] + for f in sys.argv[1:]: + data = np.load(f)['triplets'][:,1:4] + data[:,0] = 1/data[:,0] + data[:,1] = 1/data[:,1] + data2 = np.vstack((data[:,1], data[:,0], data[:,2])).transpose() + data = np.vstack((data, data2)) + alldata.append(data) + alldata = np.vstack(alldata) + #alldata = alldata[:100000] + + bandwidths = [0.001, 0.001, 1.0] + + clusterer = ManualClusterer(alldata, bandwidths=bandwidths, n_maxima=500) + clusterer.select_triplets() + diff --git a/xfel/small_cell/command_line/index2.py b/xfel/small_cell/command_line/index2.py new file mode 100644 index 00000000000..ed572e79c4f --- /dev/null +++ b/xfel/small_cell/command_line/index2.py @@ -0,0 +1,2640 @@ +import sys +import numpy as np +from dataclasses import dataclass +from typing import List, Tuple, Optional +import copy +import matplotlib.pyplot as plt +from matplotlib.colors import LinearSegmentedColormap +from tqdm import tqdm +from numba import njit +from scipy.optimize import minimize_scalar +from scipy.optimize import minimize +import itertools +from cluster2 import ManualClusterer +from cctbx import uctbx, crystal +from cctbx.sgtbx.lattice_symmetry import metric_subgroups + +class SpotPair: + + def __init__(self, q1: float, q2: float, theta: float, preserve_order=False): + """Just q1, q2, theta. Theta is in radians.""" + # Ensure q1 <= q2 + if not preserve_order and q1 > q2: + q1, q2 = q2, q1 + self.q1 = q1 + self.q2 = q2 + self.theta = theta + + def area(self) -> float: + """Calculate twice the area of the triangle formed by the vectors.""" + v1 = np.array([self.q1, 0.0]) + v2 = self.q2 * np.array([np.cos(self.theta), np.sin(self.theta)]) + return abs(np.cross(v1, v2)) + + + +class VectorPairMatch: + @classmethod + def from_pairs(cls, gen_pair, obs_pair): + return + +class PairMatch2d(VectorPairMatch): + def __init__(self, hkl1, q1, hkl2, q2, theta_rad): + self.hkl1 = hkl1 + self.hkl2 = hkl2 + self.q1 = q1 + self.q2 = q2 + self.theta_rad = theta_rad + self.is_outlier = False + +class OneVectorMatch(VectorPairMatch): + """ + q1 is the indexed vector and hkl1 is the corresponding indices in the sublattice. + """ + def __init__(self, hkl1, q1, q2, theta_rad): + self.hkl1 = hkl1 + self.q1 = q1 + self.q2 = q2 + self.theta_rad = theta_rad + + +# A couple helper functions + +def angle_between(v1, v2): + v1_u = v1/np.linalg.norm(v1) + v2_u = v2/np.linalg.norm(v2) + return np.degrees(np.arccos(np.clip(np.dot(v1_u, v2_u), -1.0, 1.0))) + +def third_vector(q, v1, v2, theta1, theta2): + """ + Compute a vector v3 with length q that forms angles theta1 and theta2 + with vectors v1 and v2 respectively. + + Parameters: + q (float): Desired length of the output vector + v1 (array-like): First reference vector + v2 (array-like): Second reference vector + theta1 (float): Desired angle with v1 (in radians) + theta2 (float): Desired angle with v2 (in radians) + + Returns: + numpy.ndarray: The computed vector v3 + """ + # Convert inputs to numpy arrays and normalize vectors + v1 = np.array(v1, dtype=float) + v2 = np.array(v2, dtype=float) + v1_norm = np.linalg.norm(v1) + v2_norm = np.linalg.norm(v2) + v1 = v1 / v1_norm + v2 = v2 / v2_norm + + # The vector we're looking for can be written as a linear combination + # of v1, v2, and their cross product: v3 = a*v1 + b*v2 + c*(v1 × v2) + + # First, get the cross product and normalize it + v1xv2 = np.cross(v1, v2) + if np.allclose(v1xv2, 0): + raise ValueError("Input vectors are parallel, solution is not unique") + v1xv2 = v1xv2 / np.linalg.norm(v1xv2) + + # The conditions are: + # q*cos(theta1) = a + b*cos(gamma) + # q*cos(theta2) = a*cos(gamma) + b + # a^2 + b^2 + c^2 = q^2 + # where gamma is the angle between v1 and v2 + + cos_gamma = np.dot(v1, v2) + sin_gamma = np.sqrt(1 - cos_gamma**2) + + # Solve for a and b + A = np.array([[1, cos_gamma], + [cos_gamma, 1]]) + b = q * np.array([np.cos(theta1), np.cos(theta2)]) + + try: + a, b = np.linalg.solve(A, b) + + # Now solve for c using the Pythagorean theorem + c_sq = q**2 - (a**2 + b**2 + 2*a*b*cos_gamma) + if c_sq < 0: + raise ValueError("No solution exists for these angles") + c = np.sqrt(c_sq) + + # There are two possible solutions (±c) + # We'll return the positive c solution + v3 = a*v1 + b*v2 + c*v1xv2 + + return v3 + + except np.linalg.LinAlgError: + raise ValueError("No solution exists for these angles") + + +def find_best_third_vector(v3_candidates: Tuple[np.ndarray, np.ndarray], + sub_basis: np.ndarray) -> np.ndarray: + """From two possible v3 positions, generate all lattice-equivalent points + and choose the closest to origin. + + Args: + v3_candidates: Two possible positions for v3 + sub_basis: Current 2x2 sublattice basis + + Returns: + The best choice for the third basis vector + """ + a, b = (np.hstack((x,0)) for x in sub_basis) + best_v3 = None + min_length = float('inf') + + for v3 in v3_candidates: + # Generate all equivalent points v3 + ha + kb + # Check a generous range of h,k values + for h in range(-3, 4): + for k in range(-3, 4): + v3_equiv = v3 + h*a + k*b + length = np.linalg.norm(v3_equiv) + #print(h, k, round(length,5)) + + if length < min_length: + min_length = length + best_v3 = v3_equiv + + return best_v3 + + +class Basis: + def __init__(self, vectors: np.ndarray, qmax: float, + q_tolerance: float = 0.001, theta_tol_degrees: float = 1.0, + symmetry=None, centering=None): + self.vectors = vectors + self.qmax = qmax + self.q_tolerance = q_tolerance + self.theta_tolerance = np.radians(theta_tol_degrees) + self.points = None + self.point_indices = None + self.pairs = None + self.indexed_pairs = [] + self.symmetry = symmetry + self.centering = centering or 'P' + + @classmethod + def from_vectors(cls, vectors: np.ndarray, reduce=True, qmax:float = 0.5, q_tol: float = 0.001, + theta_tol_deg: float = 1.0, symmetry=None, centering=None): + assert vectors.shape in ((2,2), (3,3)) + if vectors.shape == (2,2): + temp_basis = Basis2d(vectors, qmax, q_tol, theta_tol_deg) + + if not reduce: # Short circuit + return temp_basis + + temp_basis.generate_points_and_pairs_fast() + points = temp_basis.points + + reduced_vectors = [] + distances = np.linalg.norm(points, axis=1) + sort_idx = np.argsort(distances) + for i in sort_idx: + p = points[i] + if len(reduced_vectors) == 0 and np.linalg.norm(p) > 1e-6: + reduced_vectors.append(p) + elif len(reduced_vectors) == 1: + cross_prod = np.cross(reduced_vectors[0], p) + if abs(cross_prod) > 1e-6: + reduced_vectors.append(p) + break + reduced_vectors = np.vstack(reduced_vectors) + return Basis2d(reduced_vectors, qmax, q_tol, theta_tol_deg) + + elif vectors.shape == (3,3): + temp_basis = Basis3d(vectors, qmax, q_tol, theta_tol_deg, symmetry, centering) + + if not reduce: # Short circuit + return temp_basis + temp_basis.generate_points_and_pairs_fast() + points = temp_basis.points + reduced_vectors = [] + distances = np.linalg.norm(points, axis=1) + sort_idx = np.argsort(distances) + + for i in sort_idx: + p = points[i] + if len(reduced_vectors) == 0 and np.linalg.norm(p) > 1e-6: + reduced_vectors.append(p) + elif len(reduced_vectors) == 1: + cross_prod = np.cross(reduced_vectors[0], p) + if np.linalg.norm(cross_prod) > 1e-6: + reduced_vectors.append(p) + elif len(reduced_vectors) == 2: + v1, v2 = reduced_vectors + det = np.dot(np.cross(v1, v2), p) + if abs(det) > 1e-6: + reduced_vectors.append(p) + break + reduced_vectors = np.vstack(reduced_vectors) + return Basis3d(reduced_vectors, qmax, q_tol, theta_tol_deg, symmetry, centering) + + @classmethod + def from_crystal_symmetry(cls, cs, **init_kwargs): + # Make cell vectors in conventional orientation + q1, q2, q3, alpha_star, beta_star, gamma_star = cs.unit_cell().reciprocal_parameters() + alpha_star_rad = np.radians(alpha_star) + beta_star_rad = np.radians(beta_star) + gamma_star_rad = np.radians(gamma_star) + # Create basis vectors using crystallographic conventions + # First vector along x + v1 = np.array([q1, 0.0, 0.0]) + + # Second vector in xy plane + v2 = q2 * np.array([np.cos(gamma_star_rad), + np.sin(gamma_star_rad), + 0.0]) + + # Third vector using all angles + cx = np.cos(beta_star_rad) + cy = (np.cos(alpha_star_rad) - + np.cos(beta_star_rad)*np.cos(gamma_star_rad))/np.sin(gamma_star_rad) + cz = np.sqrt(1.0 - cx*cx - cy*cy) + v3 = q3 * np.array([cx, cy, cz]) + + vectors = np.vstack([v1, v2, v3]) + cr_system = cs.space_group().crystal_system() + centering = cs.space_group_info().symbol_and_number()[0] + result = cls.from_vectors(vectors, reduce=False, symmetry=cr_system, + centering=centering, **init_kwargs) + return result + + + + @classmethod + def from_params(cls, *params, reduce=True, **init_kwargs): + """ + params: q1 (A-1), q2, ga* (degrees) or q1, q2, q3, al*, be*, ga*. + reduce: if True, return the setting from the shortest suitable (non-collinear + etc) vectors. If False, return the setting as given. + """ + assert len(params) in (3,6) + + if len(params) == 3: + # 2D case: q1, q2, gamma* + q1, q2, gamma_star_rad = params + + # Create basis vectors + v1 = np.array([q1, 0.0]) + v2 = q2 * np.array([np.cos(gamma_star_rad), np.sin(gamma_star_rad)]) + vectors = np.vstack([v1, v2]) + + else: + # 3D case: q1, q2, q3, alpha*, beta*, gamma* + q1, q2, q3, alpha_star, beta_star, gamma_star = params + # Convert angles to radians + alpha_star_rad = np.radians(alpha_star) + beta_star_rad = np.radians(beta_star) + gamma_star_rad = np.radians(gamma_star) + + # Create basis vectors using crystallographic conventions + # First vector along x + v1 = np.array([q1, 0.0, 0.0]) + + # Second vector in xy plane + v2 = q2 * np.array([np.cos(gamma_star_rad), + np.sin(gamma_star_rad), + 0.0]) + + # Third vector using all angles + cx = np.cos(beta_star_rad) + cy = (np.cos(alpha_star_rad) - + np.cos(beta_star_rad)*np.cos(gamma_star_rad))/np.sin(gamma_star_rad) + cz = np.sqrt(1.0 - cx*cx - cy*cy) + v3 = q3 * np.array([cx, cy, cz]) + + vectors = np.vstack([v1, v2, v3]) + + if reduce: + # Create temporary basis to generate points + temp_basis = cls.from_vectors(vectors, **init_kwargs) + try: + temp_basis.generate_points_and_pairs_fast() + except Exception: + return None + points = temp_basis.points + + if len(params) == 3: + # Find two shortest non-collinear vectors + basis_vectors = [] + distances = np.linalg.norm(points, axis=1) + sort_idx = np.argsort(distances) + + for i in sort_idx: + p = points[i] + if len(basis_vectors) == 0 and np.linalg.norm(p) > 1e-6: + basis_vectors.append(p) + elif len(basis_vectors) == 1: + cross_prod = np.cross(basis_vectors[0], p) + if abs(cross_prod) > 1e-6: + basis_vectors.append(p) + break + vectors = np.vstack(basis_vectors) + + else: + # 3D reduction (might want to implement a more sophisticated method) + # For now, just find three shortest non-coplanar vectors + basis_vectors = [] + distances = np.linalg.norm(points, axis=1) + sort_idx = np.argsort(distances) + + for i in sort_idx: + p = points[i] + if len(basis_vectors) == 0 and np.linalg.norm(p) > 1e-6: + basis_vectors.append(p) + elif len(basis_vectors) == 1: + cross_prod = np.cross(basis_vectors[0], p) + if np.linalg.norm(cross_prod) > 1e-6: + basis_vectors.append(p) + elif len(basis_vectors) == 2: + v1, v2 = basis_vectors + det = np.dot(np.cross(v1, v2), p) + if abs(det) > 1e-6: + basis_vectors.append(p) + break + vectors = np.vstack(basis_vectors) + + return cls.from_vectors(vectors, **init_kwargs) + + def match(self, pair: SpotPair) -> Tuple[Optional[dict], str]: + raise NotImplementedError + + def fom_1d(self, pairs, delta=.001): + """Calculate figure of merit as percentage of observed q-values within threshold of lattice points. + + Args: + pairs: List of SpotPair objects with observed q-values + delta: Tolerance threshold in Å-1 + + Returns: + Percentage (0-100) of observed q-values that match lattice points + """ + if not hasattr(self, 'points') or self.points is None: + self.generate_points_and_pairs_fast() + + # Extract all observed q-values + q_obs = [] + for p in pairs: + q_obs.extend([p.q1, p.q2]) + q_obs = np.array(q_obs) + + # Calculate all q-values from lattice points + q_calc = np.linalg.norm(self.points, axis=1) + + # Reshape for broadcasting + q_obs_col = q_obs.reshape(-1, 1) # shape: (n_obs, 1) + q_calc_row = q_calc.reshape(1, -1) # shape: (1, n_calc) + + # Compute differences between all pairs + diffs = np.abs(q_obs_col - q_calc_row) # shape: (n_obs, n_calc) + + has_match = np.any(diffs < delta, axis=1) + + # Compute percentage + pct_matched = 100.0 * np.sum(has_match) / len(q_obs) + + return pct_matched + + + def generate_points_and_pairs_fast(self): + """Generate unique lattice point pairs preserving q1 ordering.""" + # Generate points + max_indices = np.ceil(self.qmax / np.linalg.norm(self.vectors, axis=1)) + ranges = [np.arange(-n, n+1) for n in max_indices.astype(int)] + + # Generate mesh grid based on dimension + if len(ranges) == 2: # 2D + h_range = np.concatenate([[0], np.arange(1, max_indices[0] + 1)]) + k_range = np.arange(-max_indices[1], max_indices[1] + 1) + H, K = np.meshgrid(h_range, k_range, indexing='ij') + indices = np.column_stack((H.flatten(), K.flatten())) + else: # 3D + h_range = np.concatenate([[0], np.arange(1, max_indices[0] + 1)]) + k_range = np.arange(-max_indices[1], max_indices[1] + 1) + l_range = np.arange(-max_indices[2], max_indices[2] + 1) + H, K, L = np.meshgrid(h_range, k_range, l_range, indexing='ij') + indices = np.column_stack((H.flatten(), K.flatten(), L.flatten())) + + # Filter indices by centering type + + lattice_condition_dict = { + None: lambda x: True, + 'P': lambda x: True, + 'A': lambda x: (x[1] + x[2]) % 2 == 0, + 'B': lambda x: (x[0] + x[2]) % 2 == 0, + 'C': lambda x: (x[0] + x[1]) % 2 == 0, + 'I': lambda x: (x[0] + x[1] + x[2]) % 2 == 0, + 'F': lambda x: x[0]%2 == x[1]%2 == x[2]%2, + 'hR': lambda x: (-x[0] + x[1] + x[2]) % 3 == 0, + 'R': lambda x: (-x[0] + x[1] + x[2]) % 3 == 0, + } + lattice_condition = lattice_condition_dict[self.centering] + lattice_mask = np.array([lattice_condition(row) for row in indices]) + indices = indices[lattice_mask] + + + # Generate points + points = indices @ self.vectors + + # Filter by magnitude + magnitudes = np.linalg.norm(points, axis=1) + mask = (magnitudes <= self.qmax) & (magnitudes > 0) # Exclude origin + + # Keep only points within qmax + filtered_points = points[mask] + filtered_indices = indices[mask] + filtered_magnitudes = magnitudes[mask] + self.points = filtered_points + self.point_indices = filtered_indices + + # Sort by magnitude (this will make q1_values naturally sorted) + sort_idx = np.argsort(filtered_magnitudes) + sorted_points = filtered_points[sort_idx] + sorted_indices = filtered_indices[sort_idx] + sorted_magnitudes = filtered_magnitudes[sort_idx] + + # Create inverted versions + inverted_indices = -sorted_indices + inverted_points = inverted_indices @ self.vectors + + # Create normalized vectors for dot products + n_points = len(sorted_points) + normalized_points = sorted_points / sorted_magnitudes[:, np.newaxis] + inverted_normalized = inverted_points / np.linalg.norm(inverted_points, axis=1)[:, np.newaxis] + + # Create full matrices with NaN padding outside upper triangle + + # 1. First create mask for upper triangle (excluding diagonal) + triu_mask = np.triu(np.ones((n_points, n_points)), k=1) + nan_mask = ~triu_mask.astype(bool) + + # 2. Create q matrices (sorted by construction) + q1_matrix = np.broadcast_to(sorted_magnitudes[:, np.newaxis], (n_points, n_points)) + q2_matrix = np.broadcast_to(sorted_magnitudes[np.newaxis, :], (n_points, n_points)) + + # 3. Calculate theta matrices + dot_products_normal = np.dot(normalized_points, normalized_points.T) + theta_matrix_normal = np.arccos(np.clip(dot_products_normal, -1, 1)) + + dot_products_inverted = np.dot(normalized_points, inverted_normalized.T) + theta_matrix_inverted = np.arccos(np.clip(dot_products_inverted, -1, 1)) + + # 4. Create hkl matrices + # For normal pairs + hkl1_shape = (n_points, n_points, sorted_indices.shape[1]) + hkl1_matrix_normal = np.broadcast_to(sorted_indices[:, np.newaxis, :], hkl1_shape) + hkl2_matrix_normal = np.broadcast_to(sorted_indices[np.newaxis, :, :], hkl1_shape) + + # For inverted pairs + hkl1_matrix_inverted = hkl1_matrix_normal.copy() + hkl2_matrix_inverted = np.broadcast_to(inverted_indices[np.newaxis, :, :], hkl1_shape) + + # 5. Apply NaN mask to matrices + q1_matrix_normal = q1_matrix.copy() + q1_matrix_normal[nan_mask] = np.nan + + q2_matrix_normal = q2_matrix.copy() + q2_matrix_normal[nan_mask] = np.nan + + theta_matrix_normal[nan_mask] = np.nan + + q1_matrix_inverted = q1_matrix.copy() + q1_matrix_inverted[nan_mask] = np.nan + + q2_matrix_inverted = q2_matrix.copy() + q2_matrix_inverted[nan_mask] = np.nan + + theta_matrix_inverted[nan_mask] = np.nan + + # 6. Stack normal and inverted matrices horizontally + q1_stacked = np.hstack([q1_matrix_normal, q1_matrix_inverted]) + q2_stacked = np.hstack([q2_matrix_normal, q2_matrix_inverted]) + theta_stacked = np.hstack([theta_matrix_normal, theta_matrix_inverted]) + + # 7. Flatten stacked matrices + q1_flat = q1_stacked.flatten() + q2_flat = q2_stacked.flatten() + theta_flat = theta_stacked.flatten() + + # 8. Remove NaN values + valid_mask = ~np.isnan(q1_flat) + self.q1_values = q1_flat[valid_mask] + self.q2_values = q2_flat[valid_mask] + self.theta_values = theta_flat[valid_mask] + + # Handle hkl values (more complex due to extra dimension) + hkl_dim = sorted_indices.shape[1] + + # Apply NaN mask to hkl matrices via boolean indexing + # Create "is NaN" array matching hkl shape + nan_mask_expanded = np.broadcast_to(nan_mask[:, :, np.newaxis], + (n_points, n_points, hkl_dim)) + + # Set invalid hkls to a dummy value (will be filtered out later) + hkl1_matrix_normal = np.where(nan_mask_expanded, -999, hkl1_matrix_normal) + hkl2_matrix_normal = np.where(nan_mask_expanded, -999, hkl2_matrix_normal) + hkl1_matrix_inverted = np.where(nan_mask_expanded, -999, hkl1_matrix_inverted) + hkl2_matrix_inverted = np.where(nan_mask_expanded, -999, hkl2_matrix_inverted) + + # Stack hkl matrices horizontally + hkl1_stacked = np.hstack([hkl1_matrix_normal.reshape(n_points, -1), + hkl1_matrix_inverted.reshape(n_points, -1)]) + hkl2_stacked = np.hstack([hkl2_matrix_normal.reshape(n_points, -1), + hkl2_matrix_inverted.reshape(n_points, -1)]) + + # Reshape to get original structure back + hkl1_flat = hkl1_stacked.reshape(-1, hkl_dim) + hkl2_flat = hkl2_stacked.reshape(-1, hkl_dim) + + # Apply same valid mask as for q values + self.hkl1_values = hkl1_flat[valid_mask] + self.hkl2_values = hkl2_flat[valid_mask] + + # Ensure q1 <= q2 (swap if needed) + swap_mask = self.q1_values > self.q2_values + if np.any(swap_mask): + print('swap') + self.q1_values[swap_mask], self.q2_values[swap_mask] = self.q2_values[swap_mask], self.q1_values[swap_mask].copy() + self.hkl1_values[swap_mask], self.hkl2_values[swap_mask] = self.hkl2_values[swap_mask], self.hkl1_values[swap_mask].copy() + def match_pairs(self, pairs): + self.all_pairs = [] + for p in pairs: + result, status = self.match(p) + self.all_pairs.append((p, result, status)) + + def reindex_pairs(self): + self.match_pairs([p[0] for p in self.all_pairs]) + + + + + + +class Basis2d(Basis): + + def index_percent(self): + all = 0 + hits = 0 + for p in self.all_pairs: + all += 1 + if p[2]=='indexed_2d': + hits += 1 + return hits/all + def fom_2d(self, pairs): + hits = 0 + for p in pairs: + result = self.match(p) + if result[1]=='indexed_2d': + hits += 1 + return 100 * hits/len(pairs) + def match(self, pair: SpotPair) -> Tuple[Optional[dict], str]: + + if not hasattr(self, 'q1_values'): + self.generate_points_and_pairs_fast() + + # Compute differences with input pair + dq1 = np.abs(self.q1_values - pair.q1) + dq2 = np.abs(self.q2_values - pair.q2) + dtheta = np.abs(self.theta_values - pair.theta) + + # Create mask for matches within tolerance + matches = (dq1 < self.q_tolerance) & (dq2 < self.q_tolerance) & (dtheta < self.theta_tolerance) + + if np.any(matches): + # Get the first match (or could find best one later) + idx = np.where(matches)[0][0] + hkl1 = self.hkl1_values[idx] + hkl2 = self.hkl2_values[idx] + + result = PairMatch2d(hkl1, pair.q1, hkl2, pair.q2, pair.theta) + return result, 'indexed_2d' + + qmags = np.linalg.norm(self.points, axis=1) + dq1 = np.abs(pair.q1 - qmags) + dq2 = np.abs(pair.q2 - qmags) + i_best = np.argmin(np.minimum(dq1, dq2)) + one_vec_match = False + if dq1[i_best] < self.q_tolerance: + one_vec_match = True + q1 = pair.q1 + q2 = pair.q2 + elif dq2[i_best] < self.q_tolerance: + one_vec_match = True + q1 = pair.q2 + q2 = pair.q1 + if one_vec_match: + return OneVectorMatch(self.point_indices[i_best], q1, q2, pair.theta), 'one_vector' + return None, 'unindexed' + + def flag_outliers(self, multiplier=2, q_weight=1000, theta_weight=111): + self.pairs_costs = [] + for p in self.all_pairs: + if p[2]=='indexed_2d': + self.pairs_costs.append(( + p, + self.compute_pair_cost(self.vectors, p[1], q_weight, theta_weight)) + ) + max_delta = np.median([pc[1] for pc in self.pairs_costs]) * multiplier + for p, c in self.pairs_costs: + p[1].is_outlier = c > max_delta + def plot_costs(self, q_weight=1000, theta_weight=111, cost_max=2): + costs = [pc[1] for pc in self.pairs_costs] + costs_inlier = [pc[1] for pc in self.pairs_costs if not pc[0][1].is_outlier] + costs_outlier = [pc[1] for pc in self.pairs_costs if pc[0][1].is_outlier] +# for p in self.sublattice_indexed: +# costs.append(self.sub_basis.compute_pair_cost( +# self.sub_basis.vectors, p, q_weight, theta_weight, use_outliers=True)) +# costs_inlier = [c for c,p in zip(costs, self.sublattice_indexed) if not p.is_outlier] +# costs_outlier = [c for c,p in zip(costs, self.sublattice_indexed) if p.is_outlier] + bins = np.linspace(0, cost_max, 50) + plt.hist(costs_inlier, bins=bins) + plt.hist(costs_outlier, bins=bins, color='red') + plt.show() + + + + def astar(self): + return(np.linalg.norm(self.vectors[0])) + def bstar(self): + return(np.linalg.norm(self.vectors[1])) + def gammastar(self): + result = np.degrees(np.arccos( + np.dot(self.vectors[0], self.vectors[1]) / (a * b))) + return result + + def __str__(self): + """Format 2D sublattice parameters.""" + # Calculate a, b, gamma from basis vectors + a = np.linalg.norm(self.vectors[0]) + b = np.linalg.norm(self.vectors[1]) + gamma = np.degrees(np.arccos( + np.dot(self.vectors[0], self.vectors[1]) / (a * b))) + + return f"a={a:.5f}, b={b:.5f}, gamma={gamma:.2f}°" + + def area(self): + return abs(np.cross(*self.vectors)) + + def doubled_cells(self) -> list: + """Generate all doubled and tripled variants of the current 2D basis. + + Returns: + List of 7 Basis2d objects: 3 doubled cells followed by 4 tripled cells + """ + # Get current reciprocal basis vectors + a_star = self.vectors[0].copy() + b_star = self.vectors[1].copy() + + result = [] + + # DOUBLED CELLS (3 cases) + # 1. Double a: a*' = a*/2, b*' = b* + basis1 = self.vectors.copy() + basis1[0] = a_star / 2 + result.append(basis1) + + # 2. Double b: a*' = a*, b*' = b*/2 + basis2 = self.vectors.copy() + basis2[1] = b_star / 2 + result.append(basis2) + + # 3. Double along (1,1): a*' = (a*+b*)/2, b*' = (a*-b*)/2 + basis3 = self.vectors.copy() + basis3[0] = (a_star + b_star) / 2 + basis3[1] = (a_star - b_star) / 2 + result.append(basis3) + + # TRIPLED CELLS (4 cases) + # 4. Triple a: a*' = a*/3, b*' = b* + basis4 = self.vectors.copy() + basis4[0] = a_star / 3 + result.append(basis4) + + # 5. Triple b: a*' = a*, b*' = b*/3 + basis5 = self.vectors.copy() + basis5[1] = b_star / 3 + result.append(basis5) + + # 6. Triple along (1,1): points at (1/3,1/3) and (2/3,-1/3) + basis6 = self.vectors.copy() + basis6[0] = (a_star + b_star) / 3 + basis6[1] = (2*a_star - b_star) / 3 + result.append(basis6) + + # 7. Triple along (1,-1): points at (1/3,-1/3) and (2/3,1/3) + basis7 = self.vectors.copy() + basis7[0] = (a_star - b_star) / 3 + basis7[1] = (2*a_star + b_star) / 3 + result.append(basis7) + + # Create new Basis2d objects + final_result = [] + for basis in result: + new_basis = type(self).from_vectors( + basis, + qmax=self.qmax, + q_tol=self.q_tolerance, + theta_tol_deg=np.degrees(self.theta_tolerance) + ) + final_result.append(new_basis) + + return final_result + + def compute_pair_cost(self, basis_vectors: np.ndarray, pair: PairMatch2d, + q_weight: float = 1.0, theta_weight: float = 111.0, use_outliers=False) -> float: + """Compute cost for a single indexed pair using Miller indices.""" + if pair.is_outlier and not use_outliers: return 0 + # Compute q-vectors from Miller indices + q1_calc = pair.hkl1[0]*basis_vectors[0] + pair.hkl1[1]*basis_vectors[1] + q2_calc = pair.hkl2[0]*basis_vectors[0] + pair.hkl2[1]*basis_vectors[1] + + # Compare magnitudes + q1_calc_mag = np.linalg.norm(q1_calc) + q2_calc_mag = np.linalg.norm(q2_calc) + q_error = (abs(q1_calc_mag - pair.q1) + + abs(q2_calc_mag - pair.q2)) + + # Compare angle + cos_theta_calc = np.dot(q1_calc, q2_calc) / (q1_calc_mag * q2_calc_mag) + theta_calc = np.arccos(np.clip(cos_theta_calc, -1, 1)) + theta_error = abs(theta_calc - pair.theta_rad) + + return q_weight * q_error + theta_weight * theta_error + + def cost(self, params, pairs, q_weight=1000, theta_weight=111.0): + """Total cost for given lattice parameters.""" + a_star, b_star, gamma_star = params + + # Basis vectors in standard orientation + v1 = np.array([a_star, 0.0]) + v2 = b_star * np.array([np.cos(gamma_star), np.sin(gamma_star)]) + basis_vectors = np.vstack([v1, v2]) + + # Sum costs from all pairs + return sum(self.compute_pair_cost(basis_vectors, pair, q_weight, theta_weight) + for pair in pairs) + + def plot_cost_histogram(self, pairs=None, q_weight=1000, theta_weight=111.0): + + if pairs is None: + pairs = [p[1] for p in self.all_pairs if p[2]=='indexed_2d'] + #copied from refine, sue me + a_star = np.linalg.norm(self.vectors[0]) + b_star = np.linalg.norm(self.vectors[1]) + gamma_star = np.arccos(np.dot(self.vectors[0], self.vectors[1]) / + (a_star * b_star)) + v1 = np.array([a_star, 0.0]) + v2 = b_star * np.array([np.cos(gamma_star), np.sin(gamma_star)]) + basis_vectors = np.vstack([v1, v2]) + + costs = [self.compute_pair_cost(basis_vectors, p, q_weight, theta_weight) for p in pairs] + plt.hist(costs, bins=100) + plt.show() + + def refine(self, pairs=None, q_weight: float = 1000.0, theta_weight: float = 111.0) -> None: + """Refine lattice parameters in place.""" + + if pairs is None: + pairs = [p[1] for p in self.all_pairs if p[2]=='indexed_2d'] + + # Initial parameter vector + a_star = np.linalg.norm(self.vectors[0]) + b_star = np.linalg.norm(self.vectors[1]) + gamma_star = np.arccos(np.dot(self.vectors[0], self.vectors[1]) / + (a_star * b_star)) + initial_params = np.array([a_star, b_star, gamma_star]) + ip = copy.deepcopy(initial_params) + + # Run optimization + result = minimize( + lambda p: self.cost(p, pairs, q_weight, theta_weight), + initial_params, + method='BFGS', + #options={'gtol': 1e-8} + ) + + if not result.success: + print(result.message) + + # Extract refined parameters and update basis + a_star, b_star, gamma_star = result.x + if gamma_star < np.pi/2: gamma_star = np.pi - gamma_star + self.vectors = np.array([[a_star, 0.0], + [b_star * np.cos(gamma_star), + b_star * np.sin(gamma_star)]]) + + # Regenerate points and pairs with new basis + self.generate_points_and_pairs_fast() + + + + +class Basis3d(Basis): + + def index_percent(self): + all = 0 + hits = 0 + for p in self.all_pairs: + all += 1 + if p[2]=='indexed_3d': + hits += 1 + return hits/all + + def flag_outliers(self, multiplier=2, q_weight=1000, theta_weight=111): + self.pairs_costs = [] + for p in self.all_pairs: + if p[2]=='indexed_3d': + self.pairs_costs.append(( + p, + self.compute_pair_cost(self.vectors, p[1], q_weight, theta_weight)) + ) + max_delta = np.median([pc[1] for pc in self.pairs_costs]) * multiplier + for p, c in self.pairs_costs: + p[1].is_outlier = c > max_delta + + def plot_costs(self, q_weight=1000, theta_weight=111): + costs = [pc[1] for pc in self.pairs_costs] + costs_inlier = [pc[1] for pc in self.pairs_costs if not pc[0][1].is_outlier] + costs_outlier = [pc[1] for pc in self.pairs_costs if pc[0][1].is_outlier] +# for p in self.sublattice_indexed: +# costs.append(self.sub_basis.compute_pair_cost( +# self.sub_basis.vectors, p, q_weight, theta_weight, use_outliers=True)) +# costs_inlier = [c for c,p in zip(costs, self.sublattice_indexed) if not p.is_outlier] +# costs_outlier = [c for c,p in zip(costs, self.sublattice_indexed) if p.is_outlier] + bins = np.linspace(0, 2, 50) + plt.hist(costs_inlier, bins=bins) + plt.hist(costs_outlier, bins=bins, color='red') + plt.show() + + def plot_costs_components(self, q_weight=1000, theta_weight=111): + theta_costs, q_costs = [],[] + for p in self.all_pairs: + if p[2]=='indexed_3d': + theta_costs.append( + self.compute_pair_cost(self.vectors, p[1], q_weight=0, theta_weight=theta_weight) + ) + q_costs.append( + self.compute_pair_cost(self.vectors, p[1], q_weight, theta_weight=0) + ) + bins = np.linspace(0, 2, 50) + fig, (ax1, ax2) = plt.subplots(2,1) + ax1.set_title('Costs (theta component)') + ax1.hist(theta_costs, bins=bins) + ax1.xaxis.set_visible(False) + ax2.set_title('Costs (q component)') + ax2.hist(q_costs, bins=bins) + plt.show() + + + def volume(self) -> float: + """Calculate the volume of the reciprocal space unit cell. + + Returns: + Volume in Å⁻³ + """ + return np.abs(np.linalg.det(self.vectors)) + + def __str__(self): + # Reciprocal cell parameters + a_star = np.linalg.norm(self.vectors[0]) + b_star = np.linalg.norm(self.vectors[1]) + c_star = np.linalg.norm(self.vectors[2]) + + alpha_star = np.degrees(np.arccos( + np.dot(self.vectors[1], self.vectors[2]) / (b_star * c_star))) + beta_star = np.degrees(np.arccos( + np.dot(self.vectors[0], self.vectors[2]) / (a_star * c_star))) + gamma_star = np.degrees(np.arccos( + np.dot(self.vectors[0], self.vectors[1]) / (a_star * b_star))) + + # Direct cell parameters + direct = self.compute_direct_cell_params(self.vectors) + + return \ + str(round(1/self.volume(), 2)) + ' ' + ','.join( [str(round(direct[x], 3)) for x in ['a','b','c','alpha','beta','gamma']] ) + ' ' + self.centering + + + def compute_direct_cell_params(self, recip_basis: np.ndarray) -> dict: + """Compute direct cell parameters from reciprocal space basis vectors. + + Args: + recip_basis: 3x3 matrix of reciprocal space basis vectors + + Returns: + dict with a,b,c (in Å) and alpha,beta,gamma (in degrees) + """ + # Compute reciprocal metric tensor G* = B B^T + G_star = recip_basis @ recip_basis.T + + # Invert to get direct metric tensor G = (G*)^-1 + G = np.linalg.inv(G_star) + + # Extract cell parameters + a = np.sqrt(G[0,0]) + b = np.sqrt(G[1,1]) + c = np.sqrt(G[2,2]) + + alpha = np.degrees(np.arccos(G[1,2] / (b*c))) + beta = np.degrees(np.arccos(G[0,2] / (a*c))) + gamma = np.degrees(np.arccos(G[0,1] / (a*b))) + + return { + 'a': a, + 'b': b, + 'c': c, + 'alpha': alpha, + 'beta': beta, + 'gamma': gamma + } + + def match(self, pair: SpotPair) -> Tuple[Optional[dict], str]: + """Find a matching pair in the lattice using binary search on sorted values.""" + + if not hasattr(self, 'q1_values'): + self.generate_points_and_pairs_fast() + + # Use binary search to find range of indices where q1 is within tolerance + q1_min = pair.q1 - self.q_tolerance + q1_max = pair.q1 + self.q_tolerance + + # Find indices where q1_values are in range [q1_min, q1_max] + left_idx = np.searchsorted(self.q1_values, q1_min, side='left') + right_idx = np.searchsorted(self.q1_values, q1_max, side='right') + + # If no values in range, return unindexed + if left_idx >= right_idx: + return None, 'unindexed' + + # Get the subset of potential matches + subset_slice = slice(left_idx, right_idx) + q2_subset = self.q2_values[subset_slice] + theta_subset = self.theta_values[subset_slice] + + # Check q2 and theta matches in the smaller subset + match_dq2 = (q2_subset > pair.q2 - self.q_tolerance) & \ + (q2_subset < pair.q2 + self.q_tolerance) + match_theta = (theta_subset > pair.theta - self.theta_tolerance) & \ + (theta_subset < pair.theta + self.theta_tolerance) + + matches = match_dq2 & match_theta + + if np.any(matches): + # Get the first match + match_idx = np.where(matches)[0][0] + full_idx = left_idx + match_idx # Adjust back to original array index + + hkl1 = self.hkl1_values[full_idx] + hkl2 = self.hkl2_values[full_idx] + + result = PairMatch2d(hkl1, pair.q1, hkl2, pair.q2, pair.theta) + return result, 'indexed_3d' + + return None, 'unindexed' + + def doubled_cells(self) -> list: + """Generate 6 cell-doubled variants of the current basis. + + Returns: + List of 6 Basis3d objects with doubled cells: + [0] Double a: a'=2a, b'=b, c'=c + [1] Double b: a'=a, b'=2b, c'=c + [2] Double c: a'=a, b'=b, c'=2c + [3] Double ab plane: a'=a+b, b'=a-b, c'=c + [4] Double ac plane: a'=a+c, b'=b, c'=a-c + [5] Double bc plane: a'=a, b'=b+c, c'=b-c + """ + result = [] + + # Get current reciprocal basis vectors + a_star = self.vectors[0].copy() + b_star = self.vectors[1].copy() + c_star = self.vectors[2].copy() + + # Create transformation matrices in reciprocal space + # When we double a direct space vector, the corresponding reciprocal vector is halved + + # 1. Double a: a*' = a*/2, b*' = b*, c*' = c* + basis1 = self.vectors.copy() + basis1[0] = a_star / 2 + result.append(basis1) + + # 2. Double b: a*' = a*, b*' = b*/2, c*' = c* + basis2 = self.vectors.copy() + basis2[1] = b_star / 2 + result.append(basis2) + + # 3. Double c: a*' = a*, b*' = b*, c*' = c*/2 + basis3 = self.vectors.copy() + basis3[2] = c_star / 2 + result.append(basis3) + + # 4. Double ab plane: a'=a+b, b'=a-b in direct space + # In reciprocal space: a*' = (a*+b*)/2, b*' = (a*-b*)/2, c*' = c* + basis4 = self.vectors.copy() + basis4[0] = (a_star + b_star) / 2 + basis4[1] = (a_star - b_star) / 2 + result.append(basis4) + + # 5. Double ac plane: a'=a+c, b'=b, c'=a-c in direct space + # In reciprocal space: a*' = (a*+c*)/2, b*' = b*, c*' = (a*-c*)/2 + basis5 = self.vectors.copy() + basis5[0] = (a_star + c_star) / 2 + basis5[2] = (a_star - c_star) / 2 + result.append(basis5) + + # 6. Double bc plane: a'=a, b'=b+c, c'=b-c in direct space + # In reciprocal space: a*' = a*, b*' = (b*+c*)/2, c*' = (b*-c*)/2 + basis6 = self.vectors.copy() + basis6[1] = (b_star + c_star) / 2 + basis6[2] = (b_star - c_star) / 2 + result.append(basis6) + + # 7. Body-centered (I-centered) + # In reciprocal space: a*'=a*, b*'=b*, c*'=a*/2+b*/2+c*/2 + basis7 = self.vectors.copy() + basis7[2] = (a_star + b_star + c_star) / 2 + result.append(basis7) + + # Now the tripled cells + + # Type 1: Along coordinate axes (3 cases) + # 1. Triple along a: a*' = a*/3 + basis1 = self.vectors.copy() + basis1[0] = a_star / 3 + result.append(basis1) + + # 2. Triple along b: b*' = b*/3 + basis2 = self.vectors.copy() + basis2[1] = b_star / 3 + result.append(basis2) + + # 3. Triple along c: c*' = c*/3 + basis3 = self.vectors.copy() + basis3[2] = c_star / 3 + result.append(basis3) + + # Type 2: Along face diagonals (6 cases) + # 4. Along (1,1,0): points (1/3,1/3,0) and (2/3,-1/3,0) + basis4 = self.vectors.copy() + basis4[0] = (a_star + b_star) / 3 + basis4[1] = (2*a_star - b_star) / 3 + result.append(basis4) + + # 5. Along (1,-1,0): points (1/3,-1/3,0) and (2/3,1/3,0) + basis5 = self.vectors.copy() + basis5[0] = (a_star - b_star) / 3 + basis5[1] = (2*a_star + b_star) / 3 + result.append(basis5) + + # 6. Along (1,0,1): points (1/3,0,1/3) and (2/3,0,-1/3) + basis6 = self.vectors.copy() + basis6[0] = (a_star + c_star) / 3 + basis6[2] = (2*a_star - c_star) / 3 + result.append(basis6) + + # 7. Along (1,0,-1): points (1/3,0,-1/3) and (2/3,0,1/3) + basis7 = self.vectors.copy() + basis7[0] = (a_star - c_star) / 3 + basis7[2] = (2*a_star + c_star) / 3 + result.append(basis7) + + # 8. Along (0,1,1): points (0,1/3,1/3) and (0,2/3,-1/3) + basis8 = self.vectors.copy() + basis8[1] = (b_star + c_star) / 3 + basis8[2] = (2*b_star - c_star) / 3 + result.append(basis8) + + # 9. Along (0,1,-1): points (0,1/3,-1/3) and (0,2/3,1/3) + basis9 = self.vectors.copy() + basis9[1] = (b_star - c_star) / 3 + basis9[2] = (2*b_star + c_star) / 3 + result.append(basis9) + + # Type 3: Along body diagonals (4 cases) + # Group 8 (body diagonal case 1) + basis10 = self.vectors.copy() + basis10[0] = (a_star - 2*b_star - 2*c_star) / 3 + basis10[1] = b_star + basis10[2] = c_star + result.append(basis10) + + # Group 9 (body diagonal case 2) + basis11 = self.vectors.copy() + basis11[0] = (a_star - b_star - 2*c_star) / 3 + basis11[1] = b_star + basis11[2] = c_star + result.append(basis11) + + # Group 10 (body diagonal case 3) + basis12 = self.vectors.copy() + basis12[0] = (a_star - 2*b_star - c_star) / 3 + basis12[1] = b_star + basis12[2] = c_star + result.append(basis12) + + # Group 11 (body diagonal case 4) + basis13 = self.vectors.copy() + basis13[0] = (a_star - b_star - c_star) / 3 + basis13[1] = b_star + basis13[2] = c_star + result.append(basis13) + + + # Create new Basis3d objects + final_result = [] + for basis in result: + new_basis = type(self).from_vectors( + basis, + qmax=self.qmax, + q_tol=self.q_tolerance, + theta_tol_deg=np.degrees(self.theta_tolerance) + ) + final_result.append(new_basis) + + + return final_result + + def compute_pair_cost(self, basis_vectors: np.ndarray, pair: PairMatch2d, + q_weight: float = 1.0, theta_weight: float = 111.0, use_outliers=False) -> float: + """Compute cost for a single indexed pair using Miller indices.""" + if not use_outliers and pair.is_outlier: return 0 + # Compute q-vectors from Miller indices + q1_calc = pair.hkl1[0]*basis_vectors[0] + pair.hkl1[1]*basis_vectors[1] + pair.hkl1[2]*basis_vectors[2] + q2_calc = pair.hkl2[0]*basis_vectors[0] + pair.hkl2[1]*basis_vectors[1] + pair.hkl2[2]*basis_vectors[2] + + # Compare magnitudes + q1_calc_mag = np.linalg.norm(q1_calc) + q2_calc_mag = np.linalg.norm(q2_calc) + q_error = (abs(q1_calc_mag - pair.q1) + + abs(q2_calc_mag - pair.q2)) + + # Compare angle + cos_theta_calc = np.dot(q1_calc, q2_calc) / (q1_calc_mag * q2_calc_mag) + theta_calc = np.arccos(np.clip(cos_theta_calc, -1, 1)) + theta_error = abs(theta_calc - pair.theta_rad) + + return q_weight * q_error + theta_weight * theta_error + + def vectors_from_params(self, params): + """Create basis vectors from refinement parameters.""" + a_star, b_star, c_star, alpha_star, beta_star, gamma_star = params + + # First vector along x + v1 = np.array([a_star, 0.0, 0.0]) + + # Second vector in xy plane + v2 = b_star * np.array([np.cos(gamma_star), + np.sin(gamma_star), + 0.0]) + + # Third vector using all angles + cx = np.cos(beta_star) + cy = (np.cos(alpha_star) - + np.cos(beta_star)*np.cos(gamma_star))/np.sin(gamma_star) + cz_sq = 1.0 - cx*cx - cy*cy + if cz_sq < 0: + # This happens with invalid angle combinations + cz = 0.0 + # Add large penalty to cost function + return None + else: + cz = np.sqrt(cz_sq) + v3 = c_star * np.array([cx, cy, cz]) + + return np.vstack([v1, v2, v3]) + + def cost(self, params, pairs, q_weight=1000, theta_weight=111.0): + """Total cost for given lattice parameters.""" + # Generate basis vectors from parameters + basis_vectors = self.vectors_from_params(params) + + # If invalid parameters, return large penalty + if basis_vectors is None: + return 1.0e6 + + # Sum costs from all pairs + return sum(self.compute_pair_cost(basis_vectors, pair, q_weight, theta_weight) + for pair in pairs) + + def plot_cost_histogram(self, pairs, q_weight=1000, theta_weight=111.0): + + #copied from refine, sue me + a_star = np.linalg.norm(self.vectors[0]) + b_star = np.linalg.norm(self.vectors[1]) + c_star = np.linalg.norm(self.vectors[2]) + + alpha_star = np.arccos(np.dot(self.vectors[1], self.vectors[2]) / + (b_star * c_star)) + beta_star = np.arccos(np.dot(self.vectors[0], self.vectors[2]) / + (a_star * c_star)) + gamma_star = np.arccos(np.dot(self.vectors[0], self.vectors[1]) / + (a_star * b_star)) + + current_params = np.array([a_star, b_star, c_star, + alpha_star, beta_star, gamma_star]) + current_vectors = self.vectors_from_params(current_params) + costs = [self.compute_pair_cost(current_params, p, q_weight, theta_weight) for p in pairs] + plt.hist(costs, bins=100) + plt.show() + + def refine(self, pairs=None, q_weight: float = 1000.0, theta_weight: float = 111.0, constrained=False) -> None: + """Refine lattice parameters in place.""" + + if pairs is None: + pairs = [p[1] for p in self.all_pairs if p[2]=='indexed_3d'] + + # Initial parameter vector from current basis + a_star = np.linalg.norm(self.vectors[0]) + b_star = np.linalg.norm(self.vectors[1]) + c_star = np.linalg.norm(self.vectors[2]) + + alpha_star = np.arccos(np.dot(self.vectors[1], self.vectors[2]) / + (b_star * c_star)) + beta_star = np.arccos(np.dot(self.vectors[0], self.vectors[2]) / + (a_star * c_star)) + gamma_star = np.arccos(np.dot(self.vectors[0], self.vectors[1]) / + (a_star * b_star)) + + if constrained: + assert self.symmetry in ['Triclinic', 'Monoclinic', 'Orthorhombic', + 'Tetragonal', 'Trigonal', 'Hexagonal', 'Cubic'] + + if not constrained or self.symmetry=='Triclinic': + initial_params = np.array([a_star, b_star, c_star, + alpha_star, beta_star, gamma_star]) + bounds = [ + (0.01, None), # a_star > 0 + (0.01, None), # b_star > 0 + (0.01, None), # c_star > 0 + (0.1, np.pi-0.1), # alpha_star in (0, pi) + (0.1, np.pi-0.1), # beta_star in (0, pi) + (0.1, np.pi-0.1) # gamma_star in (0, pi) + ] + def params_to_full(p): + return p + + elif self.symmetry == 'Monoclinic': + initial_params = np.array([a_star, b_star, c_star, beta_star]) + bounds = [ + (0.01, None), # a_star > 0 + (0.01, None), # b_star > 0 + (0.01, None), # c_star > 0 + (0.1, np.pi-0.1) # beta_star in (0, pi) + ] + def params_to_full(p): + a_star, b_star, c_star, beta_star = p + return np.array([a_star, b_star, c_star, np.pi/2, beta_star, np.pi/2]) + + elif self.symmetry == 'Orthorhombic': + initial_params = np.array([a_star, b_star, c_star]) + bounds = [ + (0.01, None), # a_star > 0 + (0.01, None), # b_star > 0 + (0.01, None), # c_star > 0 + ] + def params_to_full(p): + a_star, b_star, c_star = p + return np.array([a_star, b_star, c_star, np.pi/2, np.pi/2, np.pi/2]) + elif self.symmetry == 'Tetragonal': + initial_params = np.array([a_star, c_star]) + bounds = [ + (0.01, None), # a_star > 0 + (0.01, None), # b_star > 0 + ] + def params_to_full(p): + a_star, c_star = p + return np.array([a_star, a_star, c_star, np.pi/2, np.pi/2, np.pi/2]) + elif self.symmetry in ['Trigonal', 'Hexagonal']: + initial_params = np.array([a_star, c_star]) + bounds = [ + (0.01, None), # a_star > 0 + (0.01, None), # b_star > 0 + ] + def params_to_full(p): + a_star, c_star = p + return np.array([a_star, a_star, c_star, np.pi/2, np.pi/2, np.pi/3]) + elif self.symmetry == 'Cubic': + initial_params = np.array([a_star]) + bounds = [ + (0.01, None), # a_star > 0 + ] + def params_to_full(p): + a_star = p[0] + return np.array([a_star, a_star, a_star, np.pi/2, np.pi/2, np.pi/2]) + else: + raise RuntimeError('Unknown crystal system') + + def cost_wrapper(reduced_params): + full_params = params_to_full(reduced_params) + return self.cost(full_params, pairs, q_weight, theta_weight) + + # Run optimization + result = minimize( + cost_wrapper, + initial_params, + method='L-BFGS-B', # Use bounded optimization + bounds=bounds, + options={'ftol': 1e-10} + ) + + if not result.success: + print(f"Refinement warning: {result.message}") + + # Extract refined parameters and update basis + refined_params = params_to_full(result.x) + refined_vectors = self.vectors_from_params(refined_params) + + if refined_vectors is not None: + self.vectors = refined_vectors + + # Regenerate points and pairs with new basis + self.generate_points_and_pairs_fast() + + # Print refinement results + refined_vols = self.volume() + #print(f"Refinement complete. Final cost: {result.fun:.6f}") + #print(f"Reciprocal volume: {refined_vols:.6f} Å⁻³") + #print(f"Direct cell volume: {1/refined_vols:.1f} ų") + print(f"Refine done: {self}") + print(f"Volume: {1/refined_vols:.2f}") + + +class LatticeReconstruction: + def __init__(self, qmax: float, q_tolerance: float = 0.001, + theta_tol_degrees: float = 1): + self.qmax = qmax + self.q_tolerance = q_tolerance + self.theta_tol_degrees = theta_tol_degrees + + self.all_pairs = [] + self.sublattice_indexed = [] + self.sublattice_1vector = [] + self.unindexed = [] + self.current_basis = None + self.sub_basis = None + + def calculate_sublattice_area(self) -> float: + """Calculate area of current sublattice""" + if self.sub_basis is None: + return float('inf') + return np.abs(np.linalg.det(self.sub_basis)) + + def read_triplet(self, q1: float, q2: float, theta: float) -> None: + """Read a new q1,q2,theta triplet and update reconstruction.""" + pair = SpotPair(q1, q2, np.radians(theta)) # assume input theta in degrees + self.all_pairs.append(pair) + self.update() + + def store_triplet(self, q1, q2, theta): + pair = SpotPair(q1, q2, np.radians(theta)) + self.all_pairs.append(pair) + + def generate_2d_bases(self, scan_pts=11, scan_range=1, max_axis=40): + scan_range=np.radians(scan_range) + self.basis_candidates_2d = [] + for pair in self.all_pairs: + deltas = np.linspace(-scan_range, scan_range, scan_pts) + for d in deltas: + try: + b = Basis.from_params(pair.q1, pair.q2, pair.theta+d) + if 1/b.astar() < max_axis and 1/b.bstar() < max_axis: + self.basis_candidates_2d.append(b) + except Exception: + pass + + + def update(self) -> None: + """Main update method implementing the algorithm.""" + pair = self.all_pairs[-1] # Most recently read + + # Initialize first sublattice + if len(self.all_pairs) == 1: + self._initialize_first_sublattice(pair) + self.print_status() + return + + # Next, test for smaller sublattice + elif pair.area() < self.sub_basis.area(): + trial_basis = Basis.from_params(pair.q1, pair.q2, pair.theta) + print(f'Found smaller 2d basis: {trial_basis}') + print(f'From pair {pair.q1}, {pair.q2}, {np.degrees(pair.theta)}') + reproc = input('Reprocess? y/[n] ')=='y' + if reproc: + self.current_basis = None # Wipe out any 3d basis + self._initialize_first_sublattice(pair) + self.reprocess_all_pairs() + else: + pass + + # Try 3d indexing if possible + elif self.current_basis is not None: + result, status = self.current_basis.match(pair) + if status=='indexed': + self._store_result(result, status) + self.print_status() + return + + + # Finally, try matching in the current sublattice + else: + + # Process new pair + result, status = self.sub_basis.match(pair) + self._store_result(result, status) + + self.print_status() + + def _initialize_first_sublattice(self, pair: SpotPair) -> None: + """Set up initial 2D sublattice from first pair using least oblique cell.""" + self.sub_basis = Basis.from_params(pair.q1, pair.q2, pair.theta) + self.sl_from_pair = pair + + def set_sub_basis(self, basis): + self.sub_basis = basis + self.reprocess_all_pairs() + + def summarize_half_indexed_pairs(self) -> str: + """Create summary table of half-indexed pairs.""" + if not self.sublattice_1vector: + return "No half-indexed pairs found." + + lines = ["Half-indexed pairs for sublattice:", + "i | q[uni] | q[idx] | theta | err [Å⁻¹]"] + entries = [] + + for i, pair in enumerate(self.sublattice_1vector): + # For each pair, determine which q is indexed + q_idx, q_uni = pair.q1, pair.q2 + + # Find nearest sublattice point to the indexed vector + diffs = np.abs(np.linalg.norm(self.sub_basis.points, axis=1) - q_idx) + err = np.min(diffs) + + #lines.append(f"{i:2d} {q_uni:.4f} {q_idx:.4f} {np.degrees(pair.theta_rad):.1f} {err:.4f}") + entries.append((i, q_uni, q_idx, pair.theta_rad, err)) + + for i, q_uni, q_idx, theta_rad, err in sorted( + entries, + #key=lambda x:(round(x[1], 3), x[4]) + key=lambda x:(round(x[4], 4), round(x[1], 3)) + )[:50]: + lines.append(f"{i:2d}\t{q_uni:.4f}\t{q_idx:.4f}\t{np.degrees(theta_rad):.1f}\t{err:.4f}") + result = "\n".join(lines) + + return result + + def find_matching_pairs(self) -> List[Tuple[int, int, float]]: + """Find pairs of half-indexed vectors with matching unindexed q values. + + Returns: + List of (i1, i2, q) tuples where: + i1, i2: indices into sublattice_1vector + q: the matching q value + """ + matches = [] + n = len(self.sublattice_1vector) + + for i in range(n): + pair_i = self.sublattice_1vector[i] + # Get unindexed q value + q_i_uni = pair_i.q2 # The unindexed vector + q_i_idx = pair_i.q1 # The indexed vector + #q_i = pair_i.q2 if np.array_equal(pair_i.indexed_vector, pair_i.v1) else pair_i.q1 + + for j in range(i+1, n): + pair_j = self.sublattice_1vector[j] + # Get unindexed q value + q_j_uni = pair_j.q2 + q_j_idx = pair_j.q1 + + if abs(q_i_uni - q_j_uni) < self.q_tolerance and abs(q_i_idx - q_j_idx) > self.q_tolerance: + # Use average q value for the match + q_match = (q_i_uni + q_j_uni) / 2 + matches.append((i, j, q_match)) + + return matches + + def flag_outliers_2d(self, multiplier=2, q_weight=1000, theta_weight=111): + costs = [] + for p in self.sublattice_indexed: + costs.append(self.sub_basis.compute_pair_cost( + self.sub_basis.vectors, p, q_weight, theta_weight)) + for p, c in zip(self.sublattice_indexed, costs): + p.is_outlier = c > multiplier*np.median(costs) + def plot_costs_2d(self, q_weight=1000, theta_weight=111): + costs = [] + for p in self.sublattice_indexed: + costs.append(self.sub_basis.compute_pair_cost( + self.sub_basis.vectors, p, q_weight, theta_weight, use_outliers=True)) + costs_inlier = [c for c,p in zip(costs, self.sublattice_indexed) if not p.is_outlier] + costs_outlier = [c for c,p in zip(costs, self.sublattice_indexed) if p.is_outlier] + bins = np.linspace(0, max(costs), 50) + plt.hist(costs_inlier, bins=bins) + plt.hist(costs_outlier, bins=bins, color='red') + plt.show() + + + def reprocess_all_pairs(self) -> None: + """Clear all categorizations and reprocess all pairs except the last.""" + stored_pairs = self.all_pairs # [:-1] + self.sublattice_indexed = [] + self.sublattice_1vector = [] + self.unindexed = [] + for old_pair in stored_pairs: + result, status = self.sub_basis.match(old_pair) + self._store_result(result, status) + + def _store_result(self, result: SpotPair, status: str) -> None: + """Store a processed pair in appropriate category.""" + if status == 'indexed_2d': + self.sublattice_indexed.append(result) + elif status == 'indexed_3d': + self.indexed.append(result) + + elif status == 'one_vector': + self.sublattice_1vector.append(result) + else: # unindexed + self.unindexed.append(result) + + def generate_3d_basis_indices(self, pair1_idx, pair2_idx, delta_theta_1=None, delta_theta_2=None, inv=False): + pair1 = self.sublattice_1vector[pair1_idx] + pair2 = self.sublattice_1vector[pair2_idx] + return generate_3d_basis_pairs(pair1, pair2, delta_theta_1, delta_theta_2, inv) + + def generate_3d_basis_pairs(self, pair1, pair2, delta_theta_1=None, delta_theta_2=None, inv=False): + """Generate a 3D basis from two matching one-vector pairs.""" + + # Get indexed vectors from sublattice (2D) + v1 = np.hstack((pair1.hkl1 @ self.sub_basis.vectors, 0)) + v2 = np.hstack((pair2.hkl1 @ self.sub_basis.vectors, 0)) + + # Get common q value and both angles + q = pair1.q2 + theta1 = pair1.theta_rad + theta2 = pair2.theta_rad + + parity = -1 if inv else 1 + if delta_theta_1 is not None: + theta1 += np.radians(delta_theta_1) + if delta_theta_2 is not None: + theta2 += np.radians(delta_theta_2) + # Find possible positions for third vector (returns 3D vectors) + try: + v3_candidates = [third_vector(q, v1, parity*v2, theta1, theta2)] + except ValueError: + return None + + # Choose best third vector + v3 = find_best_third_vector(v3_candidates, self.sub_basis.vectors) + + # Create full 3D basis by extending 2D vectors + v1v2 = np.hstack((self.sub_basis.vectors, np.zeros((2,1)))) + basis_vectors = np.vstack((v1v2, v3)) + + result = Basis3d.from_vectors( + vectors=basis_vectors, + qmax=self.sub_basis.qmax, + q_tol=self.sub_basis.q_tolerance, + theta_tol_deg=np.degrees(self.sub_basis.theta_tolerance) + ) + return result + + def print_status(self, verbose=False) -> None: + """Print current status of reconstruction.""" + print(f"\nCurrent minimum sublattice: {self.sub_basis}") + if hasattr(self, 'full_basis'): + print(f"Current full lattice:\n{self.full_basis}") + else: + print("No full lattice determined yet.") + + print("\n---") + if not verbose: return + print(self.summarize_half_indexed_pairs()) + return + + # If we have matching pairs, show possible 3D cells + matches = self.find_matching_pairs() + if matches: + print("\nLattice completion possibilities:") + print("entry | i1 | i2 | cell | vol (ų) | % idx") + for i, (i1, i2, q) in enumerate(matches): + try: + basis3d = self.generate_3d_basis_indices(i1, i2) + # Calculate metrics + volume = np.abs(np.linalg.det(basis3d.vectors)) + n_indexed = len([p for p in self.all_pairs + if basis3d.match(p)[1] == 'indexed_3d']) + pct_indexed = 100 * n_indexed / len(self.all_pairs) + + print(f"{i}) {i1:2d} {i2:2d} {basis3d} {volume:.1f} {pct_indexed:.1f}") + except ValueError as e: + # Skip if no solution exists + #print(f"{i}) {i1:2d} {i2:2d} No valid solution") + pass + +# Grid search stuff, temporary + +# Define grid search parameters +def grid_search_3d_basis(recon, pair1_idx, pair2_idx, + delta_range=(-2.0, 2.0), steps=21, inv=False, + verbose=True): + # Create grid of delta theta values + delta_values = np.linspace(delta_range[0], delta_range[1], steps) + grid_shape = (steps, steps) + fom = np.zeros(grid_shape) + idx = np.zeros(grid_shape) + + # Total number of pairs + total_pairs = len(recon.all_pairs) + + if verbose: + display1_fn = display2_fn = tqdm + else: + display1_fn = lambda x: x + display2_fn = lambda x, **_: x + # Grid search + for i, delta1 in display1_fn(enumerate(delta_values)): + for j, delta2 in display2_fn(enumerate(delta_values), leave=False): + try: + # Generate 3D basis with current deltas + basis3d = recon.generate_3d_basis_pairs( + pair1_idx, pair2_idx, + delta_theta_1=delta1, + delta_theta_2=delta2, + inv=inv + ) + + # Count indexed pairs + hits = 0 + vol = basis3d.volume() + if vol > .0001: + for p in recon.all_pairs: + if basis3d.match(p)[1] != 'unindexed': + hits += 1 + + # Compute percentage + fom[i, j] = 100*hits / total_pairs * vol**(1/2) + idx[i, j] = 100*hits / total_pairs + + except Exception as e: + # Failed to generate basis (e.g., no valid solution) + fom[i, j] = 0 + idx[i, j] = 0 + + return delta_values, fom, idx + +# Run grid search for both inv=False and inv=True +def run_both_grid_searches(recon, pair1_idx, pair2_idx, + delta_range=(-2.0, 2.0), steps=21, verbose=True): + # print(f"Running grid search for pair indices {pair1_idx} and {pair2_idx}...") + # print(f"Delta range: {delta_range}, steps: {steps}") + + delta_values, fom_normal, pct_normal = grid_search_3d_basis( + recon, pair1_idx, pair2_idx, delta_range, steps, inv=False, + verbose=verbose + ) + + delta_values, fom_inv, pct_inv = grid_search_3d_basis( + recon, pair1_idx, pair2_idx, delta_range, steps, inv=True, + verbose=verbose + ) + + return delta_values, (fom_normal, pct_normal), (fom_inv, pct_inv) + +# Plot the results as heatmaps +def plot_grid_search_results(delta_values, results_normal, results_inv): + fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(12, 5)) + + # Custom colormap (white to blue) + colors = [(1, 1, 1), (0, 0, 1)] + cmap = LinearSegmentedColormap.from_list('WhiteToBlue', colors) + + # Determine shared vmax for consistent coloring + vmax = max(np.max(results_normal), np.max(results_inv)) + + # Plot normal results + im1 = ax1.imshow(results_normal, extent=[delta_values[0], delta_values[-1], + delta_values[0], delta_values[-1]], + origin='lower', cmap=cmap, vmin=0, vmax=vmax) + ax1.set_title('Normal (inv=False)') + ax1.set_xlabel('Delta theta 2 (degrees)') + ax1.set_ylabel('Delta theta 1 (degrees)') + + # Plot inverted results + im2 = ax2.imshow(results_inv, extent=[delta_values[0], delta_values[-1], + delta_values[0], delta_values[-1]], + origin='lower', cmap=cmap, vmin=0, vmax=vmax) + ax2.set_title('Inverted (inv=True)') + ax2.set_xlabel('Delta theta 2 (degrees)') + +# # Add colorbar +# cbar = fig.colorbar(im1, ax=[ax1, ax2], orientation='vertical', shrink=0.8) +# cbar.set_label('Indexing percentage (%)') + + # Add max value annotations + max_normal = np.max(results_normal) + max_normal_idx = np.unravel_index(np.argmax(results_normal), results_normal.shape) + delta1_normal = delta_values[max_normal_idx[0]] + delta2_normal = delta_values[max_normal_idx[1]] + + max_inv = np.max(results_inv) + max_inv_idx = np.unravel_index(np.argmax(results_inv), results_inv.shape) + delta1_inv = delta_values[max_inv_idx[0]] + delta2_inv = delta_values[max_inv_idx[1]] + + ax1.plot(delta2_normal, delta1_normal, 'r+', markersize=10) +# ax1.text(delta2_normal, delta1_normal, f' {max_normal:.1f}%', color='red') +# + ax2.plot(delta2_inv, delta1_inv, 'r+', markersize=10) +# ax2.text(delta2_inv, delta1_inv, f' {max_inv:.1f}%', color='red') + + plt.tight_layout() + return fig + +# Function to run everything and return the best parameters +def find_best_3d_basis(recon, pair1_idx, pair2_idx, + delta_range=(-2.0, 2.0), steps=21, plot=True, + verbose=True): + """pair1_idx and pair2_idx are actually pairs""" + # Run grid searches + delta_values, results_normal, results_inv = run_both_grid_searches( + recon, pair1_idx, pair2_idx, delta_range, steps, verbose=verbose + ) + fom_normal, pct_normal = results_normal + fom_inv, pct_inv = results_inv + + # Plot results + if plot: + fig = plot_grid_search_results(delta_values, fom_normal, fom_inv) + else: + fig = None + + # Find best parameters + max_normal = np.max(fom_normal) + max_normal_idx = np.unravel_index(np.argmax(fom_normal), fom_normal.shape) + max_normal_pct = pct_normal[max_normal_idx] + delta1_normal = delta_values[max_normal_idx[0]] + delta2_normal = delta_values[max_normal_idx[1]] + + max_inv = np.max(fom_inv) + max_inv_idx = np.unravel_index(np.argmax(fom_inv), fom_inv.shape) + max_inv_pct = pct_inv[max_inv_idx] + delta1_inv = delta_values[max_inv_idx[0]] + delta2_inv = delta_values[max_inv_idx[1]] + + # Choose best overall parameters + if max_normal >= max_inv: + best_params = { + 'delta_theta_1': delta1_normal, + 'delta_theta_2': delta2_normal, + 'inv': False, + 'fom': max_normal, + 'pct': max_normal_pct + } + else: + best_params = { + 'delta_theta_1': delta1_inv, + 'delta_theta_2': delta2_inv, + 'inv': True, + 'fom': max_inv, + 'pct': max_inv_pct + } + + # Generate the best basis + best_basis = recon.generate_3d_basis_pairs( + pair1_idx, pair2_idx, + delta_theta_1=best_params['delta_theta_1'], + delta_theta_2=best_params['delta_theta_2'], + inv=best_params['inv'] + ) + + idx_pct = best_params['fom']/best_basis.volume()**(1/2) + if verbose: + print("\nBest parameters:") + print(f"delta_theta_1 = {best_params['delta_theta_1']:.2f} degrees") + print(f"delta_theta_2 = {best_params['delta_theta_2']:.2f} degrees") + print(f"inv = {best_params['inv']}") + print(f"Indexing percentage = {idx_pct:.1f}%") + + + return best_basis, best_params, fig + + +def sb1_callback(triplet, recon): + """ + Analyze a selected triplet and return q-values to display. + """ + RANGE=2 + PTS=11 + + q1, q2, theta_degrees = triplet + theta_rad = np.radians(theta_degrees) +# deltas = np.linspace(-RANGE,RANGE,PTS) +# candidates = [] +# for d in deltas: +# test_triplet = [q1, q2, theta_rad+d*np.pi/180] +# b = Basis.from_params(*test_triplet) +# candidates.append((b, b.fom_2d(recon.all_pairs))) +# candidates.sort(key=lambda x:x[1], reverse=True) +# sb1 = candidates[0][0] + + + sb1 = Basis.from_params(q1, q2, theta_rad) + sb1.match_pairs(recon.all_pairs) + + # Refine the basis + for _ in range(5): + sb1.reindex_pairs() + sb1.flag_outliers() + sb1.refine() + print(f'{sb1}: {sb1.index_percent()}') + + # Compute q-values (norms of the basis points) + qvals = np.linalg.norm(sb1.points, axis=1) + + points_1 = np.vstack((sb1.q1_values, sb1.q2_values, sb1.theta_values*180/np.pi)).T + points_2 = np.vstack((sb1.q2_values, sb1.q1_values, sb1.theta_values*180/np.pi)).T + points = np.vstack((points_1, points_2)) + + # Filter to reasonable range for display + qvals = qvals[(qvals > 0.05) & (qvals < 0.9)] + + return qvals, points + + +# Pre-compute grid once globally +_I_VALS = np.array([-2, -2, -2, -2, -2, -1, -1, -1, -1, -1, 0, 0, 0, 0, + 1, 1, 1, 1, 1, 2, 2, 2, 2, 2]) +_J_VALS = np.array([-2, -1, 0, 1, 2, -2, -1, 0, 1, 2, -2, -1, 1, 2, + -2, -1, 0, 1, 2, -2, -1, 0, 1, 2]) +@njit +def reduce_2d_cell(a, b, gamma_deg): + """ + Find the reduced setting of a 2D unit cell (Numba JIT version). + + Parameters: + a, b: lengths of cell vectors + gamma_rad: angle between vectors (in radians) + + Returns: + (a_red, b_red, gamma_red): reduced cell parameters + """ + gamma_rad = np.radians(gamma_deg) + # Convert to Cartesian coordinates + cos_gamma = np.cos(gamma_rad) + sin_gamma = np.sin(gamma_rad) + v1 = np.array([a, 0.0]) + v2 = np.array([b * cos_gamma, b * sin_gamma]) + + for _ in range(100): + # Calculate all lattice vectors + n_vecs = len(_I_VALS) + vectors = np.zeros((n_vecs, 2)) + for i in range(n_vecs): + vectors[i, 0] = _I_VALS[i] * v1[0] + _J_VALS[i] * v2[0] + vectors[i, 1] = _I_VALS[i] * v1[1] + _J_VALS[i] * v2[1] + + # Calculate lengths + lengths = np.zeros(n_vecs) + for i in range(n_vecs): + lengths[i] = np.sqrt(vectors[i, 0]**2 + vectors[i, 1]**2) + + # Sort by length + sorted_indices = np.argsort(lengths) + + # Find two shortest non-parallel vectors + shortest = vectors[sorted_indices[0]] + + # Find first non-parallel vector + second_shortest = None + for i in range(1, n_vecs): + vec = vectors[sorted_indices[i]] + cross = abs(shortest[0] * vec[1] - shortest[1] * vec[0]) + if cross > 1e-10: + second_shortest = vec + break + + if second_shortest is None: + break + + # Calculate norms + new_a_sq = shortest[0]**2 + shortest[1]**2 + new_b_sq = second_shortest[0]**2 + second_shortest[1]**2 + new_a = np.sqrt(new_a_sq) + new_b = np.sqrt(new_b_sq) + + # Ensure a <= b + if new_a > new_b: + shortest, second_shortest = second_shortest.copy(), shortest.copy() + new_a, new_b = new_b, new_a + new_a_sq, new_b_sq = new_b_sq, new_a_sq + + # Check if better + v1_len_sq = v1[0]**2 + v1[1]**2 + v2_len_sq = v2[0]**2 + v2[1]**2 + + if new_a_sq < v1_len_sq - 1e-10 or ( + abs(new_a_sq - v1_len_sq) < 1e-10 and new_b_sq < v2_len_sq - 1e-10 + ): + v1 = shortest.copy() + v2 = second_shortest.copy() + else: + break + + # Calculate reduced parameters + a_red = np.sqrt(v1[0]**2 + v1[1]**2) + b_red = np.sqrt(v2[0]**2 + v2[1]**2) + cos_gamma_red = (v1[0] * v2[0] + v1[1] * v2[1]) / (a_red * b_red) + gamma_red = np.arccos(cos_gamma_red) + + # Ensure conventional choice + if gamma_red > np.pi / 2: + gamma_red = np.pi - gamma_red + + return a_red, b_red, np.degrees(gamma_red) + + + +def run(): + QMIN=.05 + QMAX=.45 + recon = LatticeReconstruction(qmax=QMAX) + + fn1 = sys.argv[1] + data = np.load(fn1)['triplets'][:,1:4] + data[:,0] = 1/data[:,0] + data[:,1] = 1/data[:,1] + data2 = np.vstack((data[:,1], data[:,0], data[:,2])).transpose() + data_orig = data + data = np.vstack((data, data2)) + mask = (data[:, 0] >= QMIN) & (data[:, 0] <= QMAX) & \ + (data[:, 1] >= QMIN) & (data[:, 1] <= QMAX) + data = data[mask] + #data_red = np.array([reduce_2d_cell(*item) for item in tqdm(data_orig)]) + #data_red2 = np.vstack((data_red[:,1], data_red[:,0], data_red[:,2])).transpose() + #data_red_all = np.vstack((data_red, data_red2)) + + + if len(sys.argv)==3 and False: + for l in open(sys.argv[2]): + recon.store_triplet(*[float(x) for x in l.split()]) + else: + cl_auto = ManualClusterer(data, n_maxima=500, qmin=QMIN, qmax=QMAX) + for val in cl_auto.kde_maxima: + recon.store_triplet(*val) + + assert sys.argv[2] in ['auto', 'manual'] + # Manual select 2d sub bases + if sys.argv[2] == 'manual': +# cl1 = ManualClusterer(data_red_all, n_maxima=0, sb1_callback=sb1_callback, recon=recon, qmin=QMIN, qmax=QMAX) + cl1 = ManualClusterer(data, n_maxima=0, sb1_callback=sb1_callback, recon=recon, qmin=QMIN, qmax=QMAX) + title = "select first sub-basis" + triplets = cl1.select_triplets(title) + # Lattice doubling test case: 0.153379 0.125106 40.434691 + triplets = triplets[-1:] # keep the last one + triplets[0][2] = np.radians(triplets[0][2]) + sb1 = Basis.from_params(*triplets[0]) + + else: + # Auto select 2d sub basis + recon.generate_2d_bases(scan_range=0,scan_pts=1) + print('fom1') + basis_fom = [ + (b, b.fom_1d(recon.all_pairs) ) + for b in recon.basis_candidates_2d + ] + basis_fom.sort(key=lambda x:x[1], reverse=True) + top_1d = basis_fom #[:500] + print('fom2') + basis_fom1_fom2 = [ + (b, f, b.fom_2d(recon.all_pairs)) + for b, f in top_1d + ] + basis_fom1_fom2.sort(key=lambda x:x[2], reverse=True) + for i, b in enumerate(basis_fom1_fom2[:20]): + print(i, b[0], round(b[1], 1), round(b[2], 1)) + i_sb1 = int(input('Sub-basis 1? [0]') or 0) + sb1 = basis_fom1_fom2[i_sb1] + sb1 = sb1[0] + + + sb1.match_pairs(recon.all_pairs) + for _ in range(5): + sb1.reindex_pairs() + sb1.flag_outliers() + sb1.refine() + #sb1.plot_costs() + + print(f'sb1 index percent: {sb1.index_percent()}') + # Test doubled/tripled sub-bases + i_cell = None + while i_cell != 0: + doubled_cells = sb1.doubled_cells() + all_cells = [sb1] + doubled_cells + print('Cell doubling selection:') + print('i \t%idx') + for i, c in enumerate(all_cells): + c.match_pairs(recon.all_pairs) + print(i, '\t', c.index_percent(), '\t', str(c)) + i_cell = int(input('Cell? [0]') or 0) + sb1 = all_cells[i_cell] + for _ in range(3): + sb1.reindex_pairs() + sb1.flag_outliers() + sb1.refine() + + + # Manual lattice expansion + if sys.argv[2]=='manual': + print('manual selection start') + qvals_1 = np.linalg.norm(sb1.points, axis=1) + + cl2 = ManualClusterer(data, n_maxima=0, qvals_1=qvals_1, qmin=QMIN, qmax=QMAX) + triplets = cl2.select_triplets(title="Lattice expansion") + assert len(triplets) == 2 + pairs = [SpotPair(q1,q2,np.radians(th)) for q1,q2,th in triplets] + matches = [sb1.match(p) for p in pairs] + indexed_pairs = [] + for m in matches: + assert m[1] == 'one_vector' + indexed_pairs.append(m[0]) + recon.sub_basis = sb1 + best_basis, best_params, fig = find_best_3d_basis(recon, *indexed_pairs, delta_range=(-2,2), steps=11) + plt.show() + + # Auto lattice expansion + else: + recon.sub_basis = sb1 + one_vec = [x for x in sb1.all_pairs if x[2]=='one_vector'] + one_vec.sort(key=lambda x:x[1].q2) + to_try = [] + for i in range(len(one_vec)-1): + pair1 = one_vec[i][1] + pair2 = one_vec[i+1][1] + if np.abs(pair1.q2-pair2.q2)<.001 and np.abs(pair1.q1-pair2.q1)>.005: + to_try.append((pair1, pair2)) + results = [] + for i, pair in enumerate(to_try): + try: + basis, params, _ = find_best_3d_basis( + recon, *pair, delta_range=(0,0), steps=1, plot=False, + verbose=False + ) + results.append((i, basis, params['fom'], params['pct'])) + except AttributeError: + pass + except Exception as e: + raise + print(i) + results.sort(key=lambda x:x[2], reverse=True) + for i_results, (i, basis, fom, pct) in enumerate(results[:30]): + print(i_results, '\t', pct, '\t', basis) + i_results_best = int(input('lattice: [0] ') or 0) + i_best = results[i_results_best][0] + best_basis, best_params, fig = find_best_3d_basis(recon, *to_try[i_best], delta_range=(-.5,.5), steps=11, plot=False) + plt.show() + + best_basis.match_pairs(recon.all_pairs) + for _ in range(6): + best_basis.reindex_pairs() + best_basis.flag_outliers() + best_basis.refine() + #best_basis.plot_costs() + # Test multiples of the chosen cell + i_cell = None + while i_cell != 0: + doubled_cells = best_basis.doubled_cells() + all_cells = [best_basis] + doubled_cells + print('Cell doubling selection:') + print('i \t%idx\tCell') + for i, c in enumerate(all_cells): + c.match_pairs(recon.all_pairs) + print(i, '\t', c.index_percent(), '\t', str(c)) + i_cell = int(input('Cell? [0]') or 0) + best_basis = all_cells[i_cell] + for _ in range(3): + best_basis.reindex_pairs() + best_basis.flag_outliers() + best_basis.refine() + #best_basis.plot_costs_components() + + cell_vals = list(best_basis.compute_direct_cell_params(best_basis.vectors).values()) + uc = uctbx.unit_cell(cell_vals) + cs = crystal.symmetry(unit_cell=uc, space_group='P1') + subgroups = metric_subgroups(cs, max_delta=3) + subsyms = [x['best_subsym'] for x in subgroups.result_groups] + + symmetrized_bases = [] + for i, subsym in enumerate(subsyms): + print(f"\n======= Test symmetry {i}/{len(subsyms)} =======") + constr_basis = Basis.from_crystal_symmetry(subsym) + constr_basis.match_pairs(recon.all_pairs) + for _ in range(3): + constr_basis.reindex_pairs() + constr_basis.flag_outliers() + constr_basis.refine(constrained=True) + symmetrized_bases.append((constr_basis, constr_basis.index_percent())) + + for i, b in enumerate(symmetrized_bases): print(str(i) + '\t' + str(b[0]) + ' ' + str(b[1])) + i_final = int(input('Choice: [0]') or 0) + final_basis = symmetrized_bases[i_final][0] + final_points_1 = np.vstack((final_basis.q1_values, final_basis.q2_values, final_basis.theta_values*180/np.pi)).T + final_points_2 = np.vstack((final_basis.q2_values, final_basis.q1_values, final_basis.theta_values*180/np.pi)).T + final_points = np.vstack((final_points_1, final_points_2)) + final_qvals = np.linalg.norm(final_basis.points, axis=1) + +# final_points_red = np.array([reduce_2d_cell(*item) for item in tqdm(final_points_1)]) +# final_points_red2 = np.vstack((final_points_red[:,1], final_points_red[:,0], final_points_red[:,2])).transpose() +# final_points_red_all = np.vstack((final_points_red, final_points_red2)) +# +# cl_final = ManualClusterer(data_red_all, n_maxima=0, qvals_1=final_qvals, qmin=QMIN, qmax=QMAX, points=final_points_red_all) + cl_final = ManualClusterer(data, n_maxima=0, qvals_1=final_qvals, qmin=QMIN, qmax=QMAX, points=final_points) + _ = cl_final.select_triplets() + + +if __name__=="__main__": + run() + + +def junk(): + return ''' +#def find_third_vector(v1: np.ndarray, v2: np.ndarray, +# q: float, theta1: float, theta2: float) -> Tuple[np.ndarray, np.ndarray]: +# """Find a vector v3 given its length and angles with v1 and v2. +# +# Args: +# v1, v2: Two known vectors in xy-plane (2D arrays) +# q: Length of vector to find +# theta1: Angle between v1 and v3 +# theta2: Angle between v2 and v3 +# +# Returns: +# Two possible 3D positions for v3 (above/below v1-v2 plane) +# """ +# # Convert 2D vectors to 3D +# v1_3d = np.array([v1[0], v1[1], 0.0]) +# v2_3d = np.array([v2[0], v2[1], 0.0]) +# +# # Normalize +# v1_unit = v1_3d / np.linalg.norm(v1_3d) +# v2_unit = v2_3d / np.linalg.norm(v2_3d) +# +# # Get normal to v1-v2 plane (will be along z-axis) +# n = np.cross(v1_unit, v2_unit) +# n = n / np.linalg.norm(n) # should be [0, 0, ±1] +# +# # Solve for components +# v1v2 = np.dot(v1_unit, v2_unit) +# A = np.array([[1, v1v2], [v1v2, 1]]) +# b = q * np.array([np.cos(theta1), np.cos(theta2)]) +# a, b = np.linalg.solve(A, b) +# +# # Find c from length condition +# c_sq = q*q - (a*a + b*b + 2*a*b*v1v2) +# if c_sq < 0: +# raise ValueError("No solution exists for these constraints") +# c = np.sqrt(c_sq) +# +# # Return both possible 3D positions +# v3_plus = a*v1_unit + b*v2_unit + c*n +# v3_minus = a*v1_unit + b*v2_unit - c*n +# +# return v3_plus, v3_minus + +#class IndexedSpotPair2d(SpotPair): +# def initialize_orientation(self, basis_vectors: np.ndarray) -> None: +# """Initialize orientation by computing angle between calculated and observed vectors.""" +# # Get calculated vectors +# q1_calc = self.indices1[0]*basis_vectors[0] + self.indices1[1]*basis_vectors[1] +# q2_calc = self.indices2[0]*basis_vectors[0] + self.indices2[1]*basis_vectors[1] +# +# # Compute angle of q1_calc from x-axis +# phi_calc = np.arctan2(q1_calc[1], q1_calc[0]) +# +# # Our observed q1 is along x-axis, so this is our basic rotation +# self.init_phi = phi_calc +# +# # Check if we need to flip orientation by comparing second vector +# q2_obs = self.q2 * np.array([np.cos(self.theta), np.sin(self.theta)]) +# R = np.array([[np.cos(self.init_phi), -np.sin(self.init_phi)], +# [np.sin(self.init_phi), np.cos(self.init_phi)]]) +# q2_obs_rot = R @ q2_obs +# +# # If distance is large, try flipping +# if np.linalg.norm(q2_obs_rot - q2_calc) > np.linalg.norm(q2_obs_rot + q2_calc): +# self.init_phi += np.pi +# +# # Optional: small local optimization to refine this initial guess +# result = minimize_scalar( +# lambda dphi: self.compute_cost(basis_vectors, dphi), +# bounds=(-0.1, 0.1), # small range around initial guess +# method='bounded' +# ) +# self.init_phi += result.x +# +# def compute_cost(self, basis_vectors: np.ndarray, delta_phi: float) -> float: +# """Compute distance between rotated observed and calculated positions.""" +# # Compute vectors from indices (fixed) +# q1_calc = self.indices1[0]*basis_vectors[0] + self.indices1[1]*basis_vectors[1] +# q2_calc = self.indices2[0]*basis_vectors[0] + self.indices2[1]*basis_vectors[1] +# +# # Create observed vectors in standard orientation (q1 along x) +# q1_obs = np.array([self.q1, 0.0]) +# q2_obs = self.q2 * np.array([np.cos(self.theta), np.sin(self.theta)]) +# +# # Total rotation to apply to observed vectors +# total_phi = self.init_phi + delta_phi +# R = np.array([[np.cos(total_phi), -np.sin(total_phi)], +# [np.sin(total_phi), np.cos(total_phi)]]) +# +# # Rotate observed vectors to match calculated +# q1_obs_rot = R @ q1_obs +# q2_obs_rot = R @ q2_obs +# +# return (np.linalg.norm(q1_obs_rot - q1_calc) + +# np.linalg.norm(q2_obs_rot - q2_calc)) + + # Methods from Basis2d to support the unused merge_bases approach +# def basis_in_3d(self, plus_x): +# """Make a 3x2 matrix that multiplies a point in this basis to give its +# 3d coordinates. Initially the basis should lie in the xy plane with +# the plus_x vector aligned along 1,0,0. +# +# Args: +# plus_x: A 2-element array specifying which vector in the 2D basis +# should be aligned with the positive x-axis in 3D. +# For example, [1,0] means align the first basis vector. +# +# Returns: +# A 3x2 matrix that transforms 2D basis coordinates to 3D coordinates. +# """ +# # Calculate which vector in the original basis should align with x-axis +# vector_to_align = plus_x[0] * self.vectors[0] + plus_x[1] * self.vectors[1] +# +# # Calculate the angle to rotate this vector to align with [1,0] +# norm = np.linalg.norm(vector_to_align) +# cos_angle = vector_to_align[0] / norm +# sin_angle = vector_to_align[1] / norm +# +# # Create the rotation matrix (clockwise rotation) +# # This will rotate the vector_to_align to the positive x-axis +# rotation = np.array([ +# [cos_angle, sin_angle], +# [-sin_angle, cos_angle] +# ]) +# +# # Apply rotation to both basis vectors +# rotated_vectors = rotation @ self.vectors.T +# +# # Create the 3x2 transformation matrix +# # The first two rows contain the rotated vectors +# # The third row contains zeros (since we're in the xy plane) +# result = np.zeros((3, 2)) +# result[0, :] = rotated_vectors[0, :] # x components +# result[1, :] = rotated_vectors[1, :] # y components +# +# return result +# +# def vec3d_aligned_rotated(self, point, plusx, rotx_deg): +# """Construct the aligned basis from above and compute 3d coordinates for the given point. +# Apply a rotation around the x-axis and return the transformed coordinates. +# +# Args: +# point: A 2-element array representing a point in the 2D basis +# plusx: Specifies which vector should align with the x-axis +# rotx_deg: Rotation angle around the x-axis in degrees +# +# Returns: +# A 3D point after transformation and rotation +# """ +# # Step 1: Get the 3D basis that aligns plusx with the x-axis +# basis_3d = self.basis_in_3d(plusx) +# +# # Step 2: Transform the 2D point to 3D +# point_3d = basis_3d @ np.array(point) +# +# # Step 3: Create the rotation matrix around the x-axis +# rotx_rad = np.radians(rotx_deg) +# cos_rx = np.cos(rotx_rad) +# sin_rx = np.sin(rotx_rad) +# +# # Rotation matrix around x-axis +# # [1 0 0 ] +# # [0 cos(θ) -sin(θ)] +# # [0 sin(θ) cos(θ)] +# rot_x = np.array([ +# [1, 0, 0], +# [0, cos_rx, -sin_rx], +# [0, sin_rx, cos_rx] +# ]) +# +# # Step 4: Apply the rotation +# rotated_point = rot_x @ point_3d +# +# return rotated_point +#@dataclass +#class GeneratedSpotPair(): +# hkl1: np.ndarray +# hkl2: np.ndarray +# q1: np.float64 +# q2: np.float64 +# theta: np.float64 + +# Unused functions +#def common_axis(b1, b2): +# """For two Basis2d objects, find the closest matching q-value and return +# the indices in each basis that give the corresponding value. +# """ +# qvals_1 = np.linalg.norm(b1.points, axis=1) +# qvals_2 = np.linalg.norm(b2.points, axis=1) +# delta_best = 999 +# i1_best = -1 +# i2_best = -1 +# for i2, val in enumerate(qvals_2): +# deltas = np.abs(qvals_1-val) +# i1 = np.argmin(deltas) +# if deltas[i1] < delta_best: +# i1_best = i1 +# i2_best = i2 +# delta_best = deltas[i1] +# hk1_best = b1.point_indices[i1_best] +# hk2_best = b2.point_indices[i2_best] +# print('delta_best: ', delta_best) +# print('hk1_best: ', hk1_best, qvals_1[i1_best]) +# print('hk2_best: ', hk2_best, qvals_2[i2_best]) +# return hk1_best, hk2_best +# +#def merge_bases(b1, b2, common_axes, angle_deg): +# """Merge two 2D bases into a single 3D basis. +# +# Args: +# b1: First 2D basis +# b2: Second 2D basis +# common_axes: Two 2-tuples ((h1,k1), (h2,k2)) where (h1,k1) in the first basis +# is equivalent to (h2,k2) in the second basis +# angle_deg: Rotation angle in degrees around the common axis +# +# Returns: +# A Basis3d object representing the merged 3D basis +# """ +# +# # Extract common axis vectors from the tuples +# common_axis_b1, common_axis_b2 = common_axes +# +# # Convert the first basis to 3D, aligning common axis with x-axis +# basis1_3d = b1.basis_in_3d(common_axis_b1) +# +# # For the second basis, align with x-axis +# basis2_3d_aligned = b2.basis_in_3d(common_axis_b2) +# +# # Create the full 3D basis vectors +# # First vector: common axis (aligned with x-axis) +# v1 = np.array([basis1_3d[0, 0], 0, 0]) +# +# # Second vector: from first basis, already aligned +# v2 = np.array([basis1_3d[0, 1], basis1_3d[1, 1], 0]) +# +# # Third vector: from second basis, rotated around x-axis +# # Get a vector from the second basis that's not aligned with x-axis +# # (i.e., the second column of basis2_3d_aligned) +# v3_pre_rotation = np.array([basis2_3d_aligned[0, 1], basis2_3d_aligned[1, 1], 0]) +# +# # Apply rotation around x-axis +# rot_rad = np.radians(angle_deg) +# cos_rx = np.cos(rot_rad) +# sin_rx = np.sin(rot_rad) +# +# rot_x = np.array([ +# [1, 0, 0], +# [0, cos_rx, -sin_rx], +# [0, sin_rx, cos_rx] +# ]) +# +# v3 = rot_x @ v3_pre_rotation +# +# # Combine the vectors into a 3D basis +# vectors_3d = np.vstack([v1, v2, v3]) +# +# import IPython;IPython.embed() +# # Create and return a Basis3d object +# return Basis3d(vectors_3d, qmax=max(b1.qmax, b2.qmax), +# q_tolerance=max(b1.q_tolerance, b2.q_tolerance), +# theta_tol_degrees=max(np.degrees(b1.theta_tolerance), +# np.degrees(b2.theta_tolerance))) +# +#def merge_bases_brute(b1, b2, common_axes, angle_deg): +# """Merge two 2D bases into a single 3D basis using a brute force approach. +# +# Args: +# b1: First 2D basis +# b2: Second 2D basis +# common_axes: Two 2-tuples ((h1,k1), (h2,k2)) where (h1,k1) in the first basis +# is equivalent to (h2,k2) in the second basis +# angle_deg: Rotation angle in degrees around the common axis +# +# Returns: +# A Basis3d object representing the merged 3D basis +# """ +# +# # Generate all index pairs in a grid (-4 to 4 in each dimension) +# index_grid = list(itertools.product(range(5), range(5))) +# +# # Generate 3D points from b1 +# points_3d_b1 = [] +# for hk in index_grid: +# point_3d = b1.vec3d_aligned_rotated(hk, common_axes[0], 0) +# points_3d_b1.append(point_3d) +# +# # Generate 3D points from b2 with rotation +# points_3d_b2 = [] +# for hk in index_grid: +# point_3d = b2.vec3d_aligned_rotated(hk, common_axes[1], angle_deg) +# points_3d_b2.append(point_3d) +# +# # Generate all pairwise sums, eliminating near-duplicates +# sum_vectors = [] +# for v1 in points_3d_b1: +# for v2 in points_3d_b2: +# sum_vec = v1 + v2 +# length = np.linalg.norm(sum_vec) +# +# # Skip zero vectors +# if length < 1e-2: +# continue +# +# # Check if this is a near-duplicate of an existing vector +# is_duplicate = False +# for existing_vec, _ in sum_vectors: +# if np.linalg.norm(sum_vec - existing_vec) < 0.01: # 0.01 Å⁻¹ threshold +# is_duplicate = True +# break +# +# # If not a duplicate, add it +# if not is_duplicate: +# sum_vectors.append((sum_vec, length)) +# +# # Sort by length +# sum_vectors.sort(key=lambda x: x[1]) +# +# # Find three non-coplanar vectors +# basis_vectors = [] +# for vec, _ in sum_vectors: +# if len(basis_vectors) == 0: +# basis_vectors.append(vec) +# elif len(basis_vectors) == 1: +# # Check if not collinear +# cross_prod = np.cross(basis_vectors[0], vec) +# if np.linalg.norm(cross_prod) > 1e-6: +# basis_vectors.append(vec) +# elif len(basis_vectors) == 2: +# # Check if not coplanar +# v1, v2 = basis_vectors +# det = np.dot(np.cross(v1, v2), vec) +# if abs(det) > 1e-6: +# basis_vectors.append(vec) +# break +# +# if len(basis_vectors) < 3: +# raise ValueError("Could not find three non-coplanar vectors from the combined bases") +# +# # Stack the vectors into a 3D basis +# vectors_3d = np.vstack(basis_vectors) +# +# # Create and return a Basis3d object +# return Basis3d(vectors_3d, qmax=max(b1.qmax, b2.qmax), +# q_tolerance=max(b1.q_tolerance, b2.q_tolerance), +# theta_tol_degrees=max(np.degrees(b1.theta_tolerance), +# np.degrees(b2.theta_tolerance))) + +#class PairMatch_ab(VectorPairMatch): +# """A vector pair where the first vector is indexed in the first basis and +# the second vector is indexed in the second basis.""" +# def __init__(self, q1, q2, theta_obs, b1, b2, common_axes): +# qvals_1 = np.linalg.norm(b1.points, axis=1) +# deltas_1 = np.abs(qvals_1 - q1) +# i1 = np.argmin(deltas_1) +# hk1 = b1.point_indices[i1] +# qvals_2 = np.linalg.norm(b2.points, axis=1) +# deltas_2 = np.abs(qvals_2 - q2) +# i2 = np.argmin(deltas_2) +# hk2 = b2.point_indices[i2] +# +# self.hk1 = hk1 +# self.hk2 = hk2 +# self.b1 = b1 +# self.b2 = b2 +# self.theta_obs = theta_obs +# self.common_ax1, self.common_ax2 = common_axes +# +# #check if one hkl should be inverted +# vec1 = b1.vec3d_aligned_rotated(hk1, common_axes[0], 0) +# vec2 = b2.vec3d_aligned_rotated(hk2, common_axes[1], 0) +# if np.dot(vec1, vec2) < 0: +# self.hk2 = -1 * self.hk2 +# +# def angle_error(self, rotx_deg): +# vec1 = self.b1.vec3d_aligned_rotated( +# self.hk1, self.common_ax1, 0 +# ) +# vec2 = self.b2.vec3d_aligned_rotated( +# self.hk2, self.common_ax2, rotx_deg +# ) +# theta_calc = angle_between(vec1, vec2) +# return abs(theta_calc - self.theta_obs) +# +# pass + + +# Previous attempts from the main run method + +# assert len(triplets) == 1 +# triplets[0][2] = np.radians(triplets[0][2]) +# sb2 = Basis.from_params(*triplets[0]) +# sb2.match_pairs(recon.all_pairs) +# import IPython;IPython.embed() +# for _ in range(3): +# sb2.reindex_pairs() +# sb2.flag_outliers() +# sb2.plot_costs() +# sb2.refine() +# +# +# +# # Manual select ab pairs +# qvals_1 = np.linalg.norm(sb1.points, axis=1) +# qvals_2 = np.linalg.norm(sb2.points, axis=1) +# cl = ManualClusterer(data, n_maxima=0, qvals_1=qvals_1, qvals_2=qvals_2, qmin=.1, qmax=.5) +# triplets = cl.select_triplets("select a-b pairs") +# print(triplets) +# #triplets = [[0.15333570792827658, 0.13718166412632835, 114.6026808444415], [0.1532716587572158, 0.143893922376171, 32.211055684990406], [0.15329187415603598, 0.1836433785315679, 77.80422388061861], [0.1533567548034525, 0.25795800966128574, 50.959598431305984], [0.1667772122979944, 0.1415369213997462, 65.44425336483228], [0.1668435395895667, 0.14401664386592716, 57.66669217209992], [0.16672231651340277, 0.25795198264263536, 30.3907182392463], [0.20049139390604614, 0.1415445237130603, 25.921552407884032], [0.22693602221707057, 0.2580249296609254, 108.14834482578436], [0.27737438425400857, 0.14161107283719673, 52.546032607023136], [0.2774193570666868, 0.14393247301009263, 44.26189630562257], [0.27722723290187934, 0.28328068977243553, 52.552673251938415], [0.3069234169084582, 0.14153689240822068, 89.97024400009249], [0.30674731494727275, 0.14385171337570893, 75.9191414162387], [0.3067820484586015, 0.2578167873779439, 54.49355516464851]] +# for t in triplets: +# if t[2]>90: +# t[2] = 180-t[2] +# +# +# +# common_axes = common_axis(sb1, sb2) +# all_errors = [] +# for t in triplets: +# q1, q2, th = t +# pair = PairMatch_ab(q1, q2, th, sb1, sb2, common_axes) +# errors = [pair.angle_error(x) for x in range(361)] +# all_errors.append(errors) +# for err in all_errors: +# plt.plot(range(361), err) +# plt.show() +# +# import IPython;IPython.embed() +# print('using 2d basis: ', best_sb[0], round(best_sb[1],2), round(best_sb[2],2)) +# recon.set_sub_basis(best_sb[0]) +# import IPython;IPython.embed() +# for _ in range(2): +# recon.sub_basis.refine(recon.sublattice_indexed) +# recon.reprocess_all_pairs() +# +# recon.print_status(verbose=True) +# +# from cluster2 import ManualClusterer +# fn2 = sys.argv[2] +# data = np.load(fn2)['triplets'][:,1:4] +# data[:,0] = 1/data[:,0] +# data[:,1] = 1/data[:,1] +# data2 = np.vstack((data[:,1], data[:,0], data[:,2])).transpose() +# data = np.vstack((data, data2)) +# sublattice_qvals = np.linalg.norm(recon.sub_basis.points, axis=1) +# accept = False +# while not accept: +# cl = ManualClusterer(data, n_maxima=0, mark_qvals=sublattice_qvals) +# triplets = cl.select_triplets() +# assert len(triplets)==2 +# pairs = [SpotPair(q1,q2,np.radians(th)) for q1,q2,th in triplets] +# match_results = [recon.sub_basis.match(p) for p in pairs] +# matches_1v = [] +# for m in match_results: +# assert m[1] == 'one_vector' +# matches_1v.append(m[0]) +# best_basis, best_params, fig = find_best_3d_basis(recon, *matches_1v, delta_range=(-2.0, 2.0), steps=11) +# matches = [] +# hits = 0 +# for p in recon.all_pairs: +# result = best_basis.match(p) +# if result[1]=='indexed_3d': +# hits += 1 +# matches.append(result[0]) +# +# +# import IPython;IPython.embed() +# +# +# +# +# exit() +# #recon.read_triplet(*[float(x) for x in l.split()]) +# for _ in range(2): +# recon.sub_basis.refine(recon.sublattice_indexed) +# recon.reprocess_all_pairs() +''' diff --git a/xfel/small_cell/command_line/powder_refine_geometry.py b/xfel/small_cell/command_line/powder_refine_geometry.py new file mode 100644 index 00000000000..37a51140823 --- /dev/null +++ b/xfel/small_cell/command_line/powder_refine_geometry.py @@ -0,0 +1,158 @@ +# LIBTBX_SET_DISPATCHER_NAME cctbx.xfel.powder_refine_geometry +from __future__ import division +import logging + +from iotbx.phil import parse +from dials.util import log +from dials.util import show_mail_on_error +from dials.util.options import ArgumentParser +from xfel.small_cell.geometry_refiner import PowderGeometryRefiner + + +logger = logging.getLogger("dials.command_line.powder_refine_geometry") + +help_message = """ +Refine detector geometry using powder diffraction d-spacings. + +Examples of usage: + +# Basic usage with defaults (LaB6) +$ cctbx.xfel.powder_refine_geometry spots.expt spots.refl + +# Custom d-spacings (e.g., for Si standard) +$ cctbx.xfel.powder_refine_geometry spots.expt spots.refl \\ + reference_d_spacings=3.135,1.920,1.637,1.357 + +# Or use unit cell and space group instead +$ cctbx.xfel.powder_refine_geometry spots.expt spots.refl \\ + unit_cell=4.156,4.156,4.156,90,90,90 space_group=Pm-3m + +# Refine only XY shift (fix distance and tilt) +$ cctbx.xfel.powder_refine_geometry spots.expt spots.refl \\ + refine.distance=False refine.tilt=False + +# Specify output file +$ cctbx.xfel.powder_refine_geometry spots.expt spots.refl \\ + output.experiments=calibrated.expt + +This tool uses spotfinding output from a powder standard with known d-spacings +to refine detector geometry. It optimizes a 5-parameter detector model: +- distance: shift along detector normal +- shift1/shift2: XY shifts along fast/slow axes +- tau2/tau3: tilts around fast/slow axes + +The refinement minimizes the sum of squared differences between observed +d-spacings and the nearest reference d-spacing. + +Default reference d-spacings are for LaB6 (SRM 660): + 4.156 A (100), 2.939 A (110), 2.399 A (111), 2.078 A (200), 1.858 A (210) +""" + +phil_scope = parse( + """ +reference_d_spacings = 4.156 2.939 2.399 2.078 1.858 + .type = floats + .help = Reference d-spacings in Angstroms. Default: first 5 LaB6 peaks. \ + Either use this OR specify unit_cell and space_group. + +unit_cell = None + .type = unit_cell + .help = Unit cell to generate reference d-spacings (use with space_group). \ + If specified, this overrides reference_d_spacings. + +space_group = None + .type = space_group + .help = Space group to generate reference d-spacings (use with unit_cell). \ + If specified, this overrides reference_d_spacings. + +d_min = 1.5 + .type = float + .help = Minimum d-spacing to include in refinement + +d_max = 20 + .type = float + .help = Maximum d-spacing to include in refinement + +max_distance_inv_ang = 0.002 + .type = float + .help = Maximum distance from reference d-spacing in inverse Angstroms. \ + Reflections further than this from any reference will be excluded. + +refine { + distance = True + .type = bool + .help = Refine detector distance along normal + shift = True + .type = bool + .help = Refine XY shift (shift1 and shift2) + tilt = True + .type = bool + .help = Refine detector tilts (tau2 and tau3) +} + +output { + experiments = refined.expt + .type = path + .help = Output filename for refined experiments + log = powder_refine_geometry.log + .type = path + .help = Output log file +} +""" +) + + +class Script(object): + def __init__(self): + usage = "$ cctbx.xfel.powder_refine_geometry EXPERIMENTS REFLECTIONS [options]" + self.parser = ArgumentParser( + usage=usage, + phil=phil_scope, + epilog=help_message, + check_format=False, + read_reflections=True, + read_experiments=True, + ) + + def run(self): + params, options = self.parser.parse_args(show_diff_phil=True) + + # Validate input + if len(params.input.experiments) != 1 or len(params.input.reflections) != 1: + raise ValueError("Please provide exactly one experiments file and " + "one reflections file") + + experiments = params.input.experiments[0].data + reflections = params.input.reflections[0].data + + print(f"\nLoaded {len(experiments)} experiments with " + f"{len(reflections)} reflections") + + if params.unit_cell is not None and params.space_group is not None: + print(f"Using unit_cell: {params.unit_cell}") + print(f"Using space_group: {params.space_group.symbol_and_number()}") + else: + print(f"Reference d-spacings: {params.reference_d_spacings}") + + # Create and run refiner + refiner = PowderGeometryRefiner(experiments, reflections, params) + result = refiner.run() + + if result is not None: + # Report final geometry + refiner.report_geometry_changes() + + # Save refined experiments + refined_experiments = refiner.get_refined_experiments() + print(f"\nSaving refined experiments to {params.output.experiments}") + refined_experiments.as_file(params.output.experiments) + + print("\nRefinement complete.") + else: + print("\nNo refinement performed.") + + +if __name__ == "__main__": + with show_mail_on_error(): + script = Script() + script.run() diff --git a/xfel/small_cell/geometry_refiner.py b/xfel/small_cell/geometry_refiner.py new file mode 100644 index 00000000000..b74f7fe41e6 --- /dev/null +++ b/xfel/small_cell/geometry_refiner.py @@ -0,0 +1,321 @@ +from __future__ import division +import numpy as np +import copy + +from dials.array_family import flex +from cctbx import uctbx, miller +from scitbx import matrix +from scipy.optimize import minimize + + +class PowderGeometryRefiner: + """ + Refines detector geometry using powder diffraction d-spacings. + + Uses a 5-parameter model: + - dist: distance from origin to detector plane along normal (mm) + - shift1: shift along detector fast axis (mm) + - shift2: shift along detector slow axis (mm) + - tau2: rotation around fast axis (mrad) - tilt + - tau3: rotation around slow axis (mrad) - tilt + + NOT refined: tau1 (rotation around normal) - indeterminate from powder + """ + + def __init__(self, experiments, reflections, params): + self.experiments = experiments + self.reflections = reflections + self.params = params + + # Compute reference d-spacings from unit_cell and space_group if provided + if params.unit_cell is not None and params.space_group is not None: + print(f"Computing d-spacings from unit_cell={params.unit_cell} " + f"and space_group={params.space_group.symbol_and_number()}") + + # Get beam and detector for resolution calculation + beam = experiments[0].beam + detector = experiments[0].detector + d_min = params.d_min + d_max = params.d_max + + # Average unit cell over space group + unit_cell = params.unit_cell + + # Generate Miller indices within resolution range + generator = miller.index_generator(unit_cell, params.space_group.type(), False, d_min) + indices = generator.to_array() + + # Compute d-spacings and filter by resolution range + all_spacings = unit_cell.d(indices) + spacings_in_range = flex.sorted(flex.double([d for d in all_spacings if d_min <= d <= d_max])) + + self.reference_d = np.array(spacings_in_range) + print(f"Generated {len(self.reference_d)} reference d-spacings in range " + f"[{d_min:.3f}, {d_max:.3f}] A") + else: + self.reference_d = np.array(params.reference_d_spacings) + + # Store initial detector state + self.detector = experiments[0].detector + self.initial_state = self._get_detector_state() + + # Track which parameters to refine + self.refine_distance = params.refine.distance + self.refine_shift = params.refine.shift + self.refine_tilt = params.refine.tilt + + # Build parameter vector and bounds + self._setup_parameters() + + # Prepare reflections + self._prepare_reflections() + + def _get_detector_state(self): + """Extract fast, slow, origin, and center from detector panel.""" + panel = self.detector[0] + fast = matrix.col(panel.get_fast_axis()) + slow = matrix.col(panel.get_slow_axis()) + origin = matrix.col(panel.get_origin()) + normal = fast.cross(slow) + + # Compute panel center + size = panel.get_image_size() + pixel_size = panel.get_pixel_size() + center = origin + (size[0]/2 * pixel_size[0]) * fast + (size[1]/2 * pixel_size[1]) * slow + + return { + 'fast': fast, + 'slow': slow, + 'origin': origin, + 'normal': normal, + 'center': center, + } + + def _setup_parameters(self): + """Set up parameter vector based on what is being refined.""" + # Initial values: [dist, shift1, shift2, tau2, tau3] + # All start at 0 (representing no change from initial geometry) + self.param_names = [] + self.initial_params = [] + self.bounds = [] + + if self.refine_distance: + self.param_names.append('dist') + self.initial_params.append(0.0) + self.bounds.append((-50.0, 50.0)) # +/- 50 mm + + if self.refine_shift: + self.param_names.append('shift1') + self.initial_params.append(0.0) + self.bounds.append((-10.0, 10.0)) # +/- 10 mm + + self.param_names.append('shift2') + self.initial_params.append(0.0) + self.bounds.append((-10.0, 10.0)) # +/- 10 mm + + if self.refine_tilt: + self.param_names.append('tau2') + self.initial_params.append(0.0) + self.bounds.append((-50.0, 50.0)) # +/- 50 mrad + + self.param_names.append('tau3') + self.initial_params.append(0.0) + self.bounds.append((-50.0, 50.0)) # +/- 50 mrad + + self.initial_params = np.array(self.initial_params) + + def _prepare_reflections(self): + """Filter reflections by d-range and proximity to reference d-spacings.""" + refls = self.reflections + + # Compute initial d-spacings to filter + refls.centroid_px_to_mm(self.experiments) + refls.map_centroids_to_reciprocal_space(self.experiments) + d_star_sq = flex.pow2(refls['rlp'].norms()) + refls['d'] = uctbx.d_star_sq_as_d(d_star_sq) + + # Filter by d-range + d_min = self.params.d_min + d_max = self.params.d_max + sel = (refls['d'] >= d_min) & (refls['d'] <= d_max) + refls_in_range = refls.select(sel) + + print(f"Found {len(refls_in_range)} reflections in d-range " + f"[{d_min:.3f}, {d_max:.3f}] A") + + # Filter by proximity to reference d-spacings (in inverse angstroms) + max_dist = self.params.max_distance_inv_ang + dvals = refls_in_range['d'].as_numpy_array() + dvals_inv = 1.0 / dvals + ref_d_inv = 1.0 / self.reference_d + + # For each reflection, find minimum distance to any reference + min_distances = np.min( + np.abs(dvals_inv[:, np.newaxis] - ref_d_inv[np.newaxis, :]), + axis=1 + ) + close_to_ref = min_distances <= max_dist + sel_close = flex.bool(close_to_ref.tolist()) + self.refls_filtered = refls_in_range.select(sel_close) + + print(f"Using {len(self.refls_filtered)} reflections within " + f"{max_dist:.4f} inv. A of reference d-spacings") + + # Store pixel positions for later use + self.xyzobs_px = self.refls_filtered['xyzobs.px.value'] + self.panels = self.refls_filtered['panel'] + self.ids = self.refls_filtered['id'] + + def apply_params(self, x): + """Apply parameter vector to detector geometry.""" + # Parse parameter vector + params_dict = {} + for i, name in enumerate(self.param_names): + params_dict[name] = x[i] + + # Convert to Python floats for scitbx matrix compatibility + dist = float(params_dict.get('dist', 0.0)) + shift1 = float(params_dict.get('shift1', 0.0)) + shift2 = float(params_dict.get('shift2', 0.0)) + tau2 = float(params_dict.get('tau2', 0.0)) / 1000.0 # mrad to rad + tau3 = float(params_dict.get('tau3', 0.0)) / 1000.0 # mrad to rad + + # Get initial state + d1 = self.initial_state['fast'] + d2 = self.initial_state['slow'] + dn = self.initial_state['normal'] + origin = self.initial_state['origin'] + center = self.initial_state['center'] + + # Apply rotations around the panel CENTER (not origin) + # tau2: rotation around fast axis (d1) + # tau3: rotation around slow axis (d2) + if abs(tau2) > 1e-10 or abs(tau3) > 1e-10: + # Use axis-angle rotation around the actual detector axes + # This correctly handles the left-handed detector frame + R2 = d1.axis_and_angle_as_r3_rotation_matrix(tau2, deg=False) + R3 = d2.axis_and_angle_as_r3_rotation_matrix(tau3, deg=False) + + # Combined rotation: first tau2 around fast, then tau3 around slow + R = R3 * R2 + + # Apply rotation to axes + d1_new = R * d1 + d2_new = R * d2 + + # Rotate origin around center to preserve panel center position + origin_rotated = center + R * (origin - center) + else: + d1_new = d1 + d2_new = d2 + origin_rotated = origin + + # Apply shifts and distance change (relative to initial axes) + new_origin = (origin_rotated + + shift1 * d1 + + shift2 * d2 + + dist * dn) + + # Update detector panel + panel = self.detector[0] + panel.set_frame( + d1_new.elems, + d2_new.elems, + new_origin.elems + ) + + def compute_dvals(self): + """Compute d-spacings for all reflections with current geometry.""" + # Reset mm coordinates to force recalculation + if 'xyzobs.mm.value' in self.refls_filtered: + del self.refls_filtered['xyzobs.mm.value'] + if 'rlp' in self.refls_filtered: + del self.refls_filtered['rlp'] + + self.refls_filtered.centroid_px_to_mm(self.experiments) + self.refls_filtered.map_centroids_to_reciprocal_space(self.experiments) + d_star_sq = flex.pow2(self.refls_filtered['rlp'].norms()) + self.dvals = uctbx.d_star_sq_as_d(d_star_sq) + return self.dvals + + def find_nearest_reference(self, d_obs): + """Find closest reference d-spacing.""" + distances = np.abs(self.reference_d - d_obs) + return self.reference_d[np.argmin(distances)] + + def objective(self, x): + """Compute sum of squared residuals.""" + self.apply_params(x) + dvals = self.compute_dvals() + + residuals = [] + for d_obs in dvals: + d_ref = self.find_nearest_reference(d_obs) + residuals.append(d_obs - d_ref) + + return np.sum(np.array(residuals) ** 2) + + def run(self): + """Run refinement and return results.""" + if len(self.param_names) == 0: + print("No parameters selected for refinement!") + return None + + # Compute initial objective + initial_obj = self.objective(self.initial_params) + print(f"\nInitial objective: {initial_obj:.6f}") + print(f"Initial RMS residual: {np.sqrt(initial_obj / len(self.refls_filtered)):.6f} A") + + # Run minimization using Powell method (derivative-free, more robust) + print(f"\nRefining {len(self.param_names)} parameters: {self.param_names}") + result = minimize( + self.objective, + x0=self.initial_params, + method='Powell', + options={'maxiter': 100, 'disp': True, 'xtol': 0.001, 'ftol': 0.0001} + ) + + # Apply final parameters + self.apply_params(result.x) + final_dvals = self.compute_dvals() + + # Compute final statistics + final_obj = result.fun + print(f"\nFinal objective: {final_obj:.6f}") + print(f"Final RMS residual: {np.sqrt(final_obj / len(self.refls_filtered)):.6f} A") + + # Report parameter changes + print("\nRefined parameters:") + for i, name in enumerate(self.param_names): + value = result.x[i] + unit = "mm" if name in ['dist', 'shift1', 'shift2'] else "mrad" + print(f" {name}: {value:.4f} {unit}") + + return result + + def get_refined_experiments(self): + """Return experiments with refined detector.""" + return self.experiments + + def report_geometry_changes(self): + """Print summary of geometry changes.""" + panel = self.detector[0] + new_origin = matrix.col(panel.get_origin()) + new_fast = matrix.col(panel.get_fast_axis()) + new_slow = matrix.col(panel.get_slow_axis()) + + old_origin = self.initial_state['origin'] + old_fast = self.initial_state['fast'] + old_slow = self.initial_state['slow'] + + origin_delta = new_origin - old_origin + + # Compute tilt changes from axis differences + # tau2 (tilt around fast) shows in slow axis z-component + # tau3 (tilt around slow) shows in fast axis z-component + tau2_change = np.arcsin(new_slow[2]) - np.arcsin(old_slow[2]) + tau3_change = np.arcsin(-new_fast[2]) - np.arcsin(-old_fast[2]) + + print("\nGeometry changes:") + print(f" Origin shift: ({origin_delta[0]:.4f}, {origin_delta[1]:.4f}, {origin_delta[2]:.4f}) mm") + print(f" Tilt changes: tau2={tau2_change*1000:.3f} mrad, tau3={tau3_change*1000:.3f} mrad") diff --git a/xfel/small_cell/powder_util.py b/xfel/small_cell/powder_util.py index fcdc3d85326..3eedb6de71d 100644 --- a/xfel/small_cell/powder_util.py +++ b/xfel/small_cell/powder_util.py @@ -10,6 +10,65 @@ from scitbx.math import five_number_summary from cctbx.crystal import symmetry import cctbx.miller +import itertools + +def angle(v1, v2): + """ + Compute the angle between two cartesian vectors. + + Parameters: + v1 (numpy.ndarray): The first vector. + v2 (numpy.ndarray): The second vector. + + Returns: + float: The angle between the vectors in degrees. + """ + dot_product = np.dot(v1, v2) + magnitude_v1 = np.linalg.norm(v1) + magnitude_v2 = np.linalg.norm(v2) + cos_theta = dot_product / (magnitude_v1 * magnitude_v2) + return np.degrees(np.arccos(cos_theta)) + +def create_pairwise_plots(points, labels): + # Create a figure with 3 subplots + fig, (ax1, ax2, ax3) = plt.subplots(1, 3, figsize=(15, 5)) + + # Define custom color map + color_map = { + -1: 'lightgray', # outliers in light gray + 0: 'red', + 1: 'blue', + 2: 'orange', + 3: 'green' + } + + # Convert labels to colors + colors = [color_map[label] if label in color_map else 'gray' for label in labels] + + # Define scatter plot properties + scatter_kwargs = {'c': labels, 'cmap': 'viridis', 's': 1, 'alpha': 0.6} + scatter_kwargs = {'c': colors, 's': 1, 'alpha': 0.6} + + # Plot ab + ax1.scatter(points[:, 0], points[:, 1], **scatter_kwargs) + ax1.set_xlabel('a') + ax1.set_ylabel('b') + ax1.set_title('a vs b') + + # Plot ac + ax2.scatter(points[:, 0], points[:, 2], **scatter_kwargs) + ax2.set_xlabel('a') + ax2.set_ylabel('c') + ax2.set_title('a vs c') + + # Plot bc + ax3.scatter(points[:, 1], points[:, 2], **scatter_kwargs) + ax3.set_xlabel('b') + ax3.set_ylabel('c') + ax3.set_title('b vs c') + + plt.tight_layout() + plt.show() class Spotfinder_radial_average: @@ -44,8 +103,9 @@ def _process_pixel(self, i_panel, s0, panel, xy, value): i_bin = int( n_bins * (res_inv - d_max_inv ) / (d_min_inv - d_max_inv) ) - if i_bin < 0 or i_bin >= n_bins: return - self.current_panelsums[i_panel][i_bin] += value + if 0 <= i_bin < n_bins: + self.current_panelsums[i_panel][i_bin] += value + return res def _nearest_peak(self, x, xvalues, yvalues): i = np.searchsorted(xvalues, x, side="left") @@ -100,6 +160,15 @@ def calculate(self): assert compare_detector(ref_detector, expt.detector) expt.detector = detector + if params.angle_histogram.enable: + if 's1' not in refls.keys(): + refls.centroid_px_to_mm(expts) + refls.map_centroids_to_reciprocal_space(expts) + angles_12 = [] + angles_13 = [] + angles_23 = [] + detplot_counter = 0 + for i, expt in enumerate(expts): self.current_panelsums = [ np.zeros(params.n_bins) for _ in range(self.n_panels) @@ -115,32 +184,89 @@ def calculate(self): else: self.use_current_expt = True if i % 1000 == 0: print("experiment ", i) + if self.params.n_max is not None and i>self.params.n_max: + break s0 = expt.beam.get_s0() sel = refls['id'] == i refls_sel = refls.select(sel) xyzobses = refls_sel['xyzobs.px.value'] intensities = refls_sel['intensity.sum.value'] panels = refls_sel['panel'] - shoeboxes = refls_sel['shoebox'] + if params.angle_histogram.enable: + r1max, r1min = params.angle_histogram.range1 + r2max, r2min = params.angle_histogram.range2 + r3max, r3min = params.angle_histogram.range3 + + i_r1, i_r2 ,i_r3 = [],[],[] for i_refl in range(len(refls_sel)): self.expt_count += 1 i_panel = panels[i_refl] panel = expt.detector[i_panel] - peak_height = intensities[i_refl] if params.peak_position=="xyzobs": xy = xyzobses[i_refl][0:2] if params.peak_weighting == "intensity": value = intensities[i_refl] else: value = 1 - self._process_pixel(i_panel, s0, panel, xy, value) - if params.peak_position=="shoebox": - sb = shoeboxes[i_refl] - sbpixels = zip(sb.coords(), sb.values()) - for (x,y,_), value in sbpixels: - self._process_pixel(i_panel, s0, panel, (x,y), value) + res = self._process_pixel(i_panel, s0, panel, xy, value) + if params.angle_histogram.enable and r1max > res > r1min: + i_r1.append(i_refl) + if params.angle_histogram.enable and r2max > res > r2min: + i_r2.append(i_refl) + if params.angle_histogram.enable and r3max > res > r3min: + i_r3.append(i_refl) + + if params.angle_histogram.enable and i_r1 and i_r2 and i_r3: + i_all = i_r1 + i_r2 + i_r3 + a12, a13, a23 = [],[],[] + subsel_mask = flex.bool([n in i_all for n in range(len(refls_sel))]) + subsel_mask_inv = flex.bool([not x for x in subsel_mask]) + +# subsel = refls_sel.select(subsel_mask) +# subsel_inv = refls_sel.select(subsel_mask_inv) +# xyz1 = np.array(subsel['xyzobs.mm.value']) +# xyz2 = np.array(subsel_inv['xyzobs.mm.value']) +# plt.scatter(xyz1[:,0], xyz1[:,1], c='red', s=2) +# plt.scatter(xyz2[:,0], xyz2[:,1], c='blue', s=2) +# bc = expt.detector[0].get_beam_centre(expt.beam.get_s0()) +# plt.scatter(*bc, c='k', s=5) +# plt.xlim((100,240)) +# plt.ylim((100,240)) + + s0 = np.array(expt.beam.get_s0()) + for i1, i2, i3 in itertools.product(i_r1, i_r2, i_r3): + v1 = np.array(refls_sel[i1]['s1']) - s0 + v2 = np.array(refls_sel[i2]['s1']) - s0 + v3 = np.array(refls_sel[i3]['s1']) - s0 + a12.append(angle(v1, v2)) + a13.append(angle(v1, v3)) + a23.append(angle(v2, v3)) +# print('\n----------------------') +# print('Pairwise angles:') +# headers = '1,2: 59, 121', '1,3: 30, 150', '2,3: 25, 155' +# for header, vals in zip(headers, (a12, a13, a23)): +# print() +# print(header) +# for v in vals: print(round(v, 2)) +# plt.show() + angles_12.extend(a12) + angles_13.extend(a13) + angles_23.extend(a23) +# for i1, i2 in itertools.product(i_r1, i_r2): +# v1 = np.array(refls_sel[i1]['s1']) - s0 +# v2 = np.array(refls_sel[i2]['s1']) - s0 +# angles_12.append(angle(v1, v2)) +# for i1, i2 in itertools.product(i_r1, i_r3): +# v1 = np.array(refls_sel[i1]['s1']) - s0 +# v2 = np.array(refls_sel[i2]['s1']) - s0 +# angles_13.append(angle(v1, v2)) +# for i1, i2 in itertools.product(i_r2, i_r3): +# v1 = np.array(refls_sel[i1]['s1']) - s0 +# v2 = np.array(refls_sel[i2]['s1']) - s0 +# angles_23.append(angle(v1, v2)) + for i in range(len(self.panelsums)): self.panelsums[i] = self.panelsums[i] + self.current_panelsums[i] if self.params.filter.enable and self.params.filter.select_mode=='any': @@ -159,6 +285,22 @@ def calculate(self): for i in range(len(self.panelsums)): self.antifiltered_panelsums[i] = \ self.antifiltered_panelsums[i] + self.current_panelsums[i] + if params.angle_histogram.enable: +# from sklearn.cluster import DBSCAN +# data=np.vstack((angles_12, angles_13, angles_23)).transpose() +# dbscan = DBSCAN(eps=6, min_samples=15) +# labels = dbscan.fit_predict(data) +# create_pairwise_plots(data, labels) +# fig, (ax1, ax2, ax3) = plt.subplots(1,3) +# ax1.scatter(angles_12, angles_13, s=.5) +# ax2.scatter(angles_13, angles_23, s=.5) +# ax3.scatter(angles_12, angles_23, s=.5) + fig, (ax1, ax2, ax3) = plt.subplots(3,1) + ax1.hist(angles_12, bins=180) + ax2.hist(angles_13, bins=180) + ax3.hist(angles_23, bins=180) + + plt.show() self.dvals = np.array(self.dvals) @@ -167,7 +309,7 @@ def plot(self): d_max_inv = 1/params.d_max d_min_inv = 1/params.d_min xvalues = np.linspace(d_max_inv, d_min_inv, params.n_bins) - fig, ax = plt.subplots() + fig, ax = plt.subplots(figsize=(10, 3)) ps_maxes = [max(ps) for ps in self.panelsums] ps_max = max(ps_maxes) @@ -178,6 +320,8 @@ def plot(self): yvalues = np.array(sums) plt.plot(xvalues, yvalues+0.5*i_sums*offset) elif params.filter.enable and params.filter.plot_mode=="ratio": + print('filtered: ', self.filtered_expt_count) + print('antifiltered: ', self.antifiltered_expt_count) for x in self.filtered_panelsums: x /= self.filtered_expt_count for x in self.antifiltered_panelsums: @@ -201,7 +345,7 @@ def plot(self): if params.output.xy_file: with open(params.output.xy_file, 'w') as f: for x,y in zip(xvalues, yvalues): - f.write("{:.6f}\t{}\n".format(1/x, y)) + f.write("{:.6f}\t{}\n".format(x, y)) # Now plot the predicted peak positions if requested if params.unit_cell or params.space_group: @@ -209,21 +353,25 @@ def plot(self): sym = symmetry( unit_cell=params.unit_cell, space_group=params.space_group.group() ) - hkl_list = cctbx.miller.build_set(sym, False, d_min=params.d_min) + hkl_list = cctbx.miller.build_set(sym, True, d_min=params.d_min) dspacings = params.unit_cell.d(hkl_list.indices()) +# for hkl, d in sorted(zip(hkl_list.indices(), dspacings), key=lambda x:x[1]): +# print('{:.3f}: {}'.format(d, hkl)) dspacings_inv = 1/dspacings - pplot_min = -.05*ps_max + pplot_min = -.05*max(yvalues) for d in dspacings_inv: plt.plot((d,d),(pplot_min,0), 'r-', linewidth=1) if params.output.plot_file: + fig.tight_layout() plt.savefig(params.output.plot_file) if params.plot.interactive and params.output.peak_file: backend_list = ["TkAgg","QtAgg"] - assert (plt.get_backend() in backend_list), """Matplotlib backend not compatible with interactive peak picking. -You can set the MPLBACKEND environment varibale to change this. -Currently supported options: %s""" %backend_list + print(plt.get_backend()) +# assert (plt.get_backend() in backend_list), """Matplotlib backend not compatible with interactive peak picking. +#You can set the MPLBACKEND environment varibale to change this. +#Currently supported options: %s""" %backend_list #If a peak list output file is specified, do interactive peak picking: with open(params.output.peak_file, 'w') as f: vertical_line = ax.axvline(color='r', lw=0.8, ls='--', x=xvalues[1]) diff --git a/xfel/small_cell/small_cell.py b/xfel/small_cell/small_cell.py index cc218f2f730..48853e745d4 100644 --- a/xfel/small_cell/small_cell.py +++ b/xfel/small_cell/small_cell.py @@ -1,6 +1,7 @@ from __future__ import absolute_import, division, print_function from six.moves import range from six.moves import zip +import copy #-*- Mode: Python; c-basic-offset: 2; indent-tabs-mode: nil; tab-width: 8 -*- # # LIBTBX_PRE_DISPATCHER_INCLUDE_SH export PHENIX_GUI_ENVIRONMENT=1 @@ -924,196 +925,223 @@ def small_cell_index_detail(experiments, reflections, horiz_phil, write_output = beam = imageset.get_beam() s0 = col(beam.get_s0()) - lattice_results = small_cell_index_lattice_detail(experiments, reflections, horiz_phil) - if not lattice_results: - return None - - max_clique_len, all_spots_len, ori, indexed = lattice_results - integrated_count = 0 - - if ori is not None and horiz_phil.small_cell.write_gnuplot_input: - write_cell(ori,beam,indexed,horiz_phil) - - indexed_hkls = flex.vec2_double() - indexed_intensities = flex.double() - indexed_sigmas = flex.double() - - if ori is not None: # ok to integrate - results = [] - buffers = [] - backgrounds = [] - indexed_hkls = flex.miller_index() - indexed_intensities = flex.double() - indexed_sigmas = flex.double() - mapped_predictions = flex.vec2_double() - mapped_panels = flex.size_t() - max_signal = flex.double() - xyzobs = flex.vec3_double() - xyzvar = flex.vec3_double() - shoeboxes = flex.shoebox() - s1 = flex.vec3_double() - bbox = flex.int6() - - raw_data = imageset[0] - if not isinstance(raw_data, tuple): - raw_data = (raw_data,) - rmsd = 0 - rmsd_n = 0 - for spot in indexed: - if spot.pred is None: continue - peakpix = [] - peakvals = [] - tmp = [] - is_bad = False - panel = detector[spot.pred_panel_id] - panel_raw_data = raw_data[spot.pred_panel_id] - for p in spot.peak_pixels: - #if is_bad_pixel(panel_raw_data,p): - # is_bad = True - # break - p = (p[0]+.5,p[1]+.5) - peakpix.append(p) - tmp.append(p) - peakvals.append(panel_raw_data[int(p[1]),int(p[0])]) - if is_bad: continue - - buffers.append(grow_by(peakpix,1)) - - tmp.extend(buffers[-1]) - backgrounds.append(grow_by(tmp,1)) - tmp.extend(backgrounds[-1]) - backgrounds[-1].extend(grow_by(tmp,1)) - - background = [] - bg_vals = [] - raw_bg_sum = 0 - for p in backgrounds[-1]: - try: - i = panel_raw_data[int(p[1]),int(p[0])] - except IndexError: - continue - if i is not None and i > 0: - background.append(p) - bg_vals.append(i) - raw_bg_sum += i - - ret = reject_background_outliers(background, bg_vals) - if ret is None: - print("Not enough background pixels to integrate spot %d"%spot.ID) + if horiz_phil.indexing.stills.reflection_subsampling.enable: + all_reflections = copy.deepcopy(reflections) + subsets = range( + horiz_phil.indexing.stills.reflection_subsampling.step_start, + horiz_phil.indexing.stills.reflection_subsampling.step_stop + - horiz_phil.indexing.stills.reflection_subsampling.step_size, + -horiz_phil.indexing.stills.reflection_subsampling.step_size, + ) + else: + subsets = [100] + for i_subset, pct in enumerate(subsets): + if pct != 100: + reflections = all_reflections.select( + flex.random_permutation(len(all_reflections)) + )[: int(len(all_reflections) * pct / 100)] + + + try: + lattice_results = small_cell_index_lattice_detail(experiments, reflections, horiz_phil) + if not lattice_results: continue - background, bg_vals = ret - backgrounds[-1] = background - - bp_a,bp_b,bp_c = get_background_plane_parameters(bg_vals, background) - - intensity = 0 - bg_peak = 0 - for v,p in zip(peakvals,peakpix): - intensity += v - (bp_a*p[0] + bp_b*p[1] + bp_c) - bg_peak += bp_a*p[0] + bp_b*p[1] + bp_c - - gain = panel.get_gain() - sigma = math.sqrt(gain * (intensity + bg_peak + ((len(peakvals)/len(bg_vals))**2) * raw_bg_sum)) - - print("ID: %3d, ohkl: %s, ahkl: %s, I: %9.1f, sigI: %9.1f, RDiff: %9.6f"%( \ - spot.ID, spot.hkl.get_ohkl_str(), spot.hkl.get_ahkl_str(), intensity, sigma, - (sqr(ori.reciprocal_matrix())*spot.hkl.ohkl - spot.xyz).length())) - max_sig = panel_raw_data[int(spot.spot_dict['xyzobs.px.value'][1]),int(spot.spot_dict['xyzobs.px.value'][0])] + max_clique_len, all_spots_len, ori, indexed = lattice_results + integrated_count = 0 + + if ori is not None and horiz_phil.small_cell.write_gnuplot_input: + write_cell(ori,beam,indexed,horiz_phil) + + indexed_hkls = flex.vec2_double() + indexed_intensities = flex.double() + indexed_sigmas = flex.double() + + if ori is not None: # ok to integrate + results = [] + buffers = [] + backgrounds = [] + indexed_hkls = flex.miller_index() + indexed_intensities = flex.double() + indexed_sigmas = flex.double() + mapped_predictions = flex.vec2_double() + mapped_panels = flex.size_t() + max_signal = flex.double() + xyzobs = flex.vec3_double() + xyzvar = flex.vec3_double() + shoeboxes = flex.shoebox() + s1 = flex.vec3_double() + bbox = flex.int6() + + raw_data = imageset[0] + if not isinstance(raw_data, tuple): + raw_data = (raw_data,) + rmsd = 0 + rmsd_n = 0 + for spot in indexed: + if spot.pred is None: continue + peakpix = [] + peakvals = [] + tmp = [] + is_bad = False + panel = detector[spot.pred_panel_id] + panel_raw_data = raw_data[spot.pred_panel_id] + for p in spot.peak_pixels: + #if is_bad_pixel(panel_raw_data,p): + # is_bad = True + # break + p = (p[0]+.5,p[1]+.5) + peakpix.append(p) + tmp.append(p) + peakvals.append(panel_raw_data[int(p[1]),int(p[0])]) + if is_bad: continue + + buffers.append(grow_by(peakpix,1)) + + tmp.extend(buffers[-1]) + backgrounds.append(grow_by(tmp,1)) + tmp.extend(backgrounds[-1]) + backgrounds[-1].extend(grow_by(tmp,1)) + + background = [] + bg_vals = [] + raw_bg_sum = 0 + for p in backgrounds[-1]: + try: + i = panel_raw_data[int(p[1]),int(p[0])] + except IndexError: + continue + if i is not None and i > 0: + background.append(p) + bg_vals.append(i) + raw_bg_sum += i + + ret = reject_background_outliers(background, bg_vals) + if ret is None: + print("Not enough background pixels to integrate spot %d"%spot.ID) + continue + background, bg_vals = ret + backgrounds[-1] = background + + bp_a,bp_b,bp_c = get_background_plane_parameters(bg_vals, background) + + intensity = 0 + bg_peak = 0 + for v,p in zip(peakvals,peakpix): + intensity += v - (bp_a*p[0] + bp_b*p[1] + bp_c) + bg_peak += bp_a*p[0] + bp_b*p[1] + bp_c + + gain = panel.get_gain() + sigma = math.sqrt(gain * (intensity + bg_peak + ((len(peakvals)/len(bg_vals))**2) * raw_bg_sum)) + + print("ID: %3d, ohkl: %s, ahkl: %s, I: %9.1f, sigI: %9.1f, RDiff: %9.6f"%( \ + spot.ID, spot.hkl.get_ohkl_str(), spot.hkl.get_ahkl_str(), intensity, sigma, + (sqr(ori.reciprocal_matrix())*spot.hkl.ohkl - spot.xyz).length())) + + max_sig = panel_raw_data[int(spot.spot_dict['xyzobs.px.value'][1]),int(spot.spot_dict['xyzobs.px.value'][0])] + + s = "Orig HKL: % 4d % 4d % 4d "%(spot.hkl.ohkl.elems) + s = s + "Asu HKL: % 4d % 4d % 4d "%(spot.hkl.ahkl.elems) + s = s + "I: % 10.1f sigI: % 8.1f I/sigI: % 8.1f "%(intensity, sigma, intensity/sigma) + s = s + "Size (pix): %3d Max pix val: %6d\n"%(len(spot.peak_pixels),max_sig) + results.append(s) + + if spot.pred is None: + mapped_predictions.append((spot.spot_dict['xyzobs.px.value'][0], spot.spot_dict['xyzobs.px.value'][1])) + mapped_panels.append(spot.spot_dict['panel']) + else: + mapped_predictions.append((spot.pred[0],spot.pred[1])) + mapped_panels.append(spot.pred_panel_id) + xyzobs.append(spot.spot_dict['xyzobs.px.value']) + xyzvar.append(spot.spot_dict['xyzobs.px.variance']) + shoeboxes.append(spot.spot_dict['shoebox']) + + indexed_hkls.append(spot.hkl.ohkl.elems) + indexed_intensities.append(intensity) + indexed_sigmas.append(sigma) + max_signal.append(max_sig) + s1.append(s0+spot.xyz) + bbox.append(spot.spot_dict['bbox']) + + if spot.pred is not None: + rmsd_n += 1 + rmsd += measure_distance(col((spot.spot_dict['xyzobs.px.value'][0],spot.spot_dict['xyzobs.px.value'][1])),col(spot.pred))**2 + + if len(results) >= horiz_phil.small_cell.min_spots_to_integrate: + # Uncomment to get a text version of the integration results + #f = open(os.path.splitext(os.path.basename(path))[0] + ".int","w") + #for line in results: + # f.write(line) + #f.close() + + if write_output: + info = dict( + xbeam = refined_bcx, + ybeam = refined_bcy, + distance = distance, + wavelength = wavelength, + pointgroup = horiz_phil.small_cell.spacegroup, + observations = [cctbx.miller.set(sym,indexed_hkls).array(indexed_intensities,indexed_sigmas)], + mapped_predictions = [mapped_predictions], + mapped_panels = [mapped_panels], + model_partialities = [None], + sa_parameters = [None], + max_signal = [max_signal], + current_orientation = [ori], + current_cb_op_to_primitive = [sgtbx.change_of_basis_op()], # identity. only support primitive lattices. + pixel_size = pixel_size, + ) + G = open("int-" + os.path.splitext(os.path.basename(path))[0] +".pickle","wb") + import pickle + pickle.dump(info,G,pickle.HIGHEST_PROTOCOL) + + crystal = ori_to_crystal(ori, horiz_phil.small_cell.spacegroup) + experiments = ExperimentListFactory.from_imageset_and_crystal(imageset, crystal) + if write_output: + experiments.as_file( + os.path.splitext(os.path.basename(path).strip())[0] + "_integrated.expt" + ) + + refls = flex.reflection_table() + refls['id'] = flex.int(len(indexed_hkls), 0) + refls['panel'] = mapped_panels + refls['intensity.sum.value'] = indexed_intensities + refls['intensity.sum.variance'] = indexed_sigmas**2 + refls['xyzobs.px.value'] = xyzobs + refls['xyzobs.px.variance'] = xyzvar + refls['miller_index'] = indexed_hkls + refls['xyzcal.px'] = flex.vec3_double(mapped_predictions.parts()[0], mapped_predictions.parts()[1], flex.double(len(mapped_predictions), 0)) + refls['shoebox'] = shoeboxes + refls['entering'] = flex.bool(len(refls), False) + refls['s1'] = s1 + refls['bbox'] = bbox + + refls.centroid_px_to_mm(experiments) + + refls.set_flags(flex.bool(len(refls), True), refls.flags.indexed) + if write_output: + refls.as_pickle(os.path.splitext(os.path.basename(path).strip())[0]+"_integrated.refl") + + print("cctbx.small_cell: integrated %d spots."%len(results), end=' ') + integrated_count = len(results) + else: + raise RuntimeError("cctbx.small_cell: not enough spots to integrate (%d)."%len(results)) - s = "Orig HKL: % 4d % 4d % 4d "%(spot.hkl.ohkl.elems) - s = s + "Asu HKL: % 4d % 4d % 4d "%(spot.hkl.ahkl.elems) - s = s + "I: % 10.1f sigI: % 8.1f I/sigI: % 8.1f "%(intensity, sigma, intensity/sigma) - s = s + "Size (pix): %3d Max pix val: %6d\n"%(len(spot.peak_pixels),max_sig) - results.append(s) + if rmsd_n > 0: + print(" RMSD: %f"%math.sqrt((1/rmsd_n)*rmsd)) + else: + print(" Cannot calculate RMSD. Not enough integrated spots or not enough clique spots near predictions.") - if spot.pred is None: - mapped_predictions.append((spot.spot_dict['xyzobs.px.value'][0], spot.spot_dict['xyzobs.px.value'][1])) - mapped_panels.append(spot.spot_dict['panel']) + print("IMAGE STATS %s: spots %5d, max clique: %5d, integrated %5d spots"%(path,all_spots_len,max_clique_len,integrated_count)) + except Exception: + if i_subset == len(subsets)-1: + raise else: - mapped_predictions.append((spot.pred[0],spot.pred[1])) - mapped_panels.append(spot.pred_panel_id) - xyzobs.append(spot.spot_dict['xyzobs.px.value']) - xyzvar.append(spot.spot_dict['xyzobs.px.variance']) - shoeboxes.append(spot.spot_dict['shoebox']) - - indexed_hkls.append(spot.hkl.ohkl.elems) - indexed_intensities.append(intensity) - indexed_sigmas.append(sigma) - max_signal.append(max_sig) - s1.append(s0+spot.xyz) - bbox.append(spot.spot_dict['bbox']) - - if spot.pred is not None: - rmsd_n += 1 - rmsd += measure_distance(col((spot.spot_dict['xyzobs.px.value'][0],spot.spot_dict['xyzobs.px.value'][1])),col(spot.pred))**2 - - if len(results) >= horiz_phil.small_cell.min_spots_to_integrate: - # Uncomment to get a text version of the integration results - #f = open(os.path.splitext(os.path.basename(path))[0] + ".int","w") - #for line in results: - # f.write(line) - #f.close() - - if write_output: - info = dict( - xbeam = refined_bcx, - ybeam = refined_bcy, - distance = distance, - wavelength = wavelength, - pointgroup = horiz_phil.small_cell.spacegroup, - observations = [cctbx.miller.set(sym,indexed_hkls).array(indexed_intensities,indexed_sigmas)], - mapped_predictions = [mapped_predictions], - mapped_panels = [mapped_panels], - model_partialities = [None], - sa_parameters = [None], - max_signal = [max_signal], - current_orientation = [ori], - current_cb_op_to_primitive = [sgtbx.change_of_basis_op()], # identity. only support primitive lattices. - pixel_size = pixel_size, - ) - G = open("int-" + os.path.splitext(os.path.basename(path))[0] +".pickle","wb") - import pickle - pickle.dump(info,G,pickle.HIGHEST_PROTOCOL) - - crystal = ori_to_crystal(ori, horiz_phil.small_cell.spacegroup) - experiments = ExperimentListFactory.from_imageset_and_crystal(imageset, crystal) - if write_output: - experiments.as_file( - os.path.splitext(os.path.basename(path).strip())[0] + "_integrated.expt" - ) - - refls = flex.reflection_table() - refls['id'] = flex.int(len(indexed_hkls), 0) - refls['panel'] = mapped_panels - refls['intensity.sum.value'] = indexed_intensities - refls['intensity.sum.variance'] = indexed_sigmas**2 - refls['xyzobs.px.value'] = xyzobs - refls['xyzobs.px.variance'] = xyzvar - refls['miller_index'] = indexed_hkls - refls['xyzcal.px'] = flex.vec3_double(mapped_predictions.parts()[0], mapped_predictions.parts()[1], flex.double(len(mapped_predictions), 0)) - refls['shoebox'] = shoeboxes - refls['entering'] = flex.bool(len(refls), False) - refls['s1'] = s1 - refls['bbox'] = bbox - - refls.centroid_px_to_mm(experiments) - - refls.set_flags(flex.bool(len(refls), True), refls.flags.indexed) - if write_output: - refls.as_pickle(os.path.splitext(os.path.basename(path).strip())[0]+"_integrated.refl") - - print("cctbx.small_cell: integrated %d spots."%len(results), end=' ') - integrated_count = len(results) - else: - raise RuntimeError("cctbx.small_cell: not enough spots to integrate (%d)."%len(results)) - - if rmsd_n > 0: - print(" RMSD: %f"%math.sqrt((1/rmsd_n)*rmsd)) + continue else: - print(" Cannot calculate RMSD. Not enough integrated spots or not enough clique spots near predictions.") + print('indexed on subset ', i_subset) + break - print("IMAGE STATS %s: spots %5d, max clique: %5d, integrated %5d spots"%(path,all_spots_len,max_clique_len,integrated_count)) return max_clique_len, experiments, refls def spots_rmsd(spots): From dd96fdd4674fee19f97b2b5e5db76c1be257702f Mon Sep 17 00:00:00 2001 From: Daniel Paley Date: Thu, 26 Mar 2026 11:08:50 -0400 Subject: [PATCH 02/15] small_cell: remove superseded cake_plot_prep script This script is no longer needed and has been superseded by other tools. Co-Authored-By: Claude Sonnet 4.5 --- .../small_cell/command_line/cake_plot_prep.py | 94 ------------------- 1 file changed, 94 deletions(-) delete mode 100644 xfel/small_cell/command_line/cake_plot_prep.py diff --git a/xfel/small_cell/command_line/cake_plot_prep.py b/xfel/small_cell/command_line/cake_plot_prep.py deleted file mode 100644 index cdce3106704..00000000000 --- a/xfel/small_cell/command_line/cake_plot_prep.py +++ /dev/null @@ -1,94 +0,0 @@ -from __future__ import division -# LIBTBX_SET_DISPATCHER_NAME cctbx.xfel.small_cell.cake_plot_prep -from dials.array_family import flex -from dxtbx.model.experiment_list import ExperimentList -from dials.array_family import flex -import sys, glob -import matplotlib.pyplot as plt - - -help_str = """ -Make a cake plot from DIALS spotfinder spots - -A cake plot is the azimuthal angle of a spot on an image vs. its resolution. -Powder rings will appear as vertical stripes, with defects in geometry -causing them to appear wavy. A cake plot is also insensitive to badly masked -regions of the detector compared to a 1d radial average as the aziumuthal -angle of a spot isn't averaged into the 1d trace. - -This script creates cake.npy which is used by -cctbx.xfel.small_cell.cake_plot. Run this script first to generate it, then -run cctbx.xfel.small_cell.cake_plot - -This script expects files named "*_strong.expt" and "_strong.refl". Supply -the former and the script will seek for the latter. - -Usage (note, wild cards are permitted, but quotes are recommended): -cctbx.xfel.small_cell.cake_plot_prep "/*_strong.expt>" - -Multiprocessing support is availible using MPI. Example: -mpirun cctbx.xfel.small_cell.cake_plot_prep "/*_strong.expt>" -""" - -def run(args): - if "-h" in args or "--help" in args: - if rank == 0: - print(help_str) - return - - filenames = [] - for arg in sys.argv[1:]: - filenames.extend(glob.glob(arg)) - if not filenames: - sys.exit("No data found") - - x, y = flex.double(), flex.double() - det = None - for fn in filenames: - print (fn) - #try: - refls = flex.reflection_table.from_file(fn.split('_strong.expt')[0] + "_strong.refl") - #except OSError: - # continue - expts = ExperimentList.from_file(fn, check_format=False) - for expt_id, expt in enumerate(expts): - subset = refls.select(expt_id == refls['id']) - if len(subset) > 200: continue - det = expt.detector - for panel_id, panel in enumerate(det): - r = subset.select(subset['panel'] == panel_id) - x_, y_, _ = r['xyzobs.px.value'].parts() - pix = panel.pixel_to_millimeter(flex.vec2_double(x_, y_)) - c = panel.get_lab_coord(pix) - x.extend(c.parts()[0]) - y.extend(c.parts()[1]) - - if det: - z = flex.double(len(x), sum([p.get_origin()[2] for p in det])/len(det)) - coords = flex.vec3_double(x,y,z) - two_theta = coords.angle((0,0,-1)) - d = expts[0].beam.get_wavelength() / 2 / flex.sin(two_theta/2) - azi = flex.vec3_double(x, y, flex.double(len(x), 0)).angle((0,1,0), deg=True) - azi.set_selected(x < 0, 180+(180-azi.select(x<0))) - else: - d = flex.double() - azi = flex.double() - - import numpy as np - fig, axes = plt.subplots(1, 1, figsize=(6, 3)) - axes.plot( - 1/d.as_numpy_array(), azi.as_numpy_array(), - marker='.', linestyle='none', - markersize=0.5, alpha=0.5 - ) - axes.set_ylabel('Azimuthal Angle') - axes.set_xlabel(r'Resolution ($\mathrm{\AA}$)') - x_ticks = np.array([10, 5, 2, 1]) - axes.set_xticks(1/x_ticks) - axes.set_xticklabels(x_ticks) - fig.tight_layout() - plt.show() - -if __name__ == "__main__": - run(sys.argv[1:]) - From 7ac5fbc69332964d18dab24df8b4a508ac49ff02 Mon Sep 17 00:00:00 2001 From: Daniel Paley Date: Thu, 26 Mar 2026 15:44:55 -0400 Subject: [PATCH 03/15] Move cluster2 and index2 to smx_findexer branch These experimental indexing and clustering algorithms are specific to the smx_findexer workflow and have been moved to that dedicated branch. The smx_small_cell_processing branch now contains only general small-cell features: - geometry_refiner.py: Geometry refinement functionality - powder_refine_geometry.py: Powder-based geometry refinement tool - powder_util.py: General powder pattern utilities - small_cell.py: Core small-cell processing with reflection subsampling The experimental clustering/indexing algorithms (cluster2.py, index2.py) are now on the smx_findexer branch along with the smx_statistics worker that supports them. Co-Authored-By: Claude Sonnet 4.5 --- xfel/small_cell/command_line/cluster2.py | 1182 ---------- xfel/small_cell/command_line/index2.py | 2640 ---------------------- 2 files changed, 3822 deletions(-) delete mode 100644 xfel/small_cell/command_line/cluster2.py delete mode 100644 xfel/small_cell/command_line/index2.py diff --git a/xfel/small_cell/command_line/cluster2.py b/xfel/small_cell/command_line/cluster2.py deleted file mode 100644 index 4371312e1e6..00000000000 --- a/xfel/small_cell/command_line/cluster2.py +++ /dev/null @@ -1,1182 +0,0 @@ -import numpy as np -import matplotlib.pyplot as plt -from matplotlib.widgets import SpanSelector, RectangleSelector -from matplotlib.patches import Rectangle -from scipy.stats import gaussian_kde -from scipy.ndimage import maximum_filter -from scipy.ndimage import generate_binary_structure -from sklearn.neighbors import KernelDensity -import sys -from scipy.optimize import minimize - -import matplotlib -matplotlib.use('TkAgg') - -def fit_gaussian_peak(points, initial_center, bandwidths): - """Fit 3D Gaussian to refine peak location""" - from scipy.optimize import minimize - - def negative_log_likelihood(params): - center = params[:3] - # Gaussian kernel with anisotropic bandwidths - diff = (points - center) / bandwidths - distances_sq = np.sum(diff**2, axis=1) - log_prob = -0.5 * distances_sq - return -np.sum(log_prob) # Negative for minimization - - result = minimize( - negative_log_likelihood, - initial_center, - method='BFGS' - ) - - if result.success: - return result.x - else: - return initial_center # Fall back to grid position - -def has_saddle_between(peak1, peak2, hist_smooth, edges, bin_width, prominence_threshold=0.5): - """Check if there's a significant saddle between two peaks""" - # Sample points along the line between peaks - n_samples = 50 - t = np.linspace(0, 1, n_samples) - line_points = peak1[np.newaxis, :] * (1 - t[:, np.newaxis]) + peak2[np.newaxis, :] * t[:, np.newaxis] - - # Convert to bin indices - bin_indices = [] - for dim in range(3): - indices = np.searchsorted(edges[dim], line_points[:, dim]) - 1 - indices = np.clip(indices, 0, hist_smooth.shape[dim] - 1) - bin_indices.append(indices) - - # Get density along the line - line_densities = hist_smooth[bin_indices[0], bin_indices[1], bin_indices[2]] - - # Find minimum along path - min_density = np.min(line_densities) - - # Get densities at the two peaks - peak1_idx = [np.searchsorted(edges[i], peak1[i]) - 1 for i in range(3)] - peak2_idx = [np.searchsorted(edges[i], peak2[i]) - 1 for i in range(3)] - peak1_idx = [np.clip(idx, 0, hist_smooth.shape[i] - 1) for i, idx in enumerate(peak1_idx)] - peak2_idx = [np.clip(idx, 0, hist_smooth.shape[i] - 1) for i, idx in enumerate(peak2_idx)] - - density1 = hist_smooth[tuple(peak1_idx)] - density2 = hist_smooth[tuple(peak2_idx)] - - # Check prominence: saddle should be significantly lower than both peaks - lower_peak = min(density1, density2) - prominence = (lower_peak - min_density) / lower_peak - - result = prominence > prominence_threshold - return result - - -class ManualClusterer: - def __init__(self, data, bandwidths=[0.001, 0.001, 1.0], n_maxima=500, - qvals_1=None, qvals_2=None, sb1_callback=None, recon=None, - qmin=.1, qmax=.5, points=None, n_shortest=500): - if len(data) > 1000000: - indices = np.random.choice(len(data), size=1000000, replace=False) - data = data[indices] - self.data = data - self.bandwidths = np.array(bandwidths) - self.n_maxima = n_maxima - self.n_shortest = n_shortest - self.qvals_1 = qvals_1 - self.qvals_2 = qvals_2 - self.sb1_qvals = None - self.qmin = qmin - self.qmax = qmax - self.points = points - - self.current_selection = None - self.final_selection = None - self.means = None - self.span_patch = None - self.rect_patch = None - self.selected_triplets = [] - - self.sb1_callback = sb1_callback - self.recon = recon - - # Compute KDE maxima - if n_maxima > 0: - self.compute_pf_maxima() - self.points = self.kde_maxima - - - def select_qvals(self, title="Select q-value ranges"): - """ - Simple q-value selection mode using 1D histogram on dimension 0. - - Parameters: - ----------- - title : str - Title for the histogram window - - Returns: - -------- - list : Selected q-values (medians of selected ranges) - """ - # Create figure for histogram - self.fig_qval, self.ax_qval = plt.subplots(num='Q-value Selection', figsize=(16, 4)) - - # Initialize list to store selected q-values - self.selected_qvals = [] - self.qval_spans = [] # Store span patches for visualization - - # Plot histogram - self._plot_qval_histogram(title) - - # Set up span selector - self.qval_span = SpanSelector( - self.ax_qval, - self._on_qval_span_select, - 'horizontal', - useblit=True, - props=dict(alpha=0.5, facecolor='blue') - ) - - # Add done button - done_ax = self.fig_qval.add_axes([0.9, 0.01, 0.09, 0.05]) - self.qval_done_button = plt.Button(done_ax, 'Done') - self.qval_done_button.on_clicked(self._on_qval_done) - - # Add clear button - clear_ax = self.fig_qval.add_axes([0.8, 0.01, 0.09, 0.05]) - self.qval_clear_button = plt.Button(clear_ax, 'Clear Last') - self.qval_clear_button.on_clicked(self._on_qval_clear_last) - - # Connect key press events - self.fig_qval.canvas.mpl_connect('key_press_event', self._on_qval_key_press) - - # Show plot and wait for user interaction - plt.show(block=True) - - # Close figure - plt.close(self.fig_qval) - - return self.selected_qvals - - def _plot_qval_histogram(self, title): - """Plot histogram for q-value selection.""" - self.ax_qval.clear() - - # Create histogram - plot_range = .1, self.qmax - hist, edges = np.histogram(self.data[:, 0], bins=2000, range=plotrange) - self.ax_qval.plot(edges[:-1], hist) - self.ax_qval.set_title('abcd') - self.ax_qval.set_xlabel('q-value') - self.ax_qval.set_ylabel('Count') - - # Add tick marks for previously selected q-values - if self.qvals_1 is not None: - self.ax_qval.plot(self.qvals_1, np.zeros_like(self.qvals_1), - '|', color='blue', markersize=25, markeredgewidth=2, label='qvals_1') - if self.qvals_2 is not None: - self.ax_qval.plot(self.qvals_2, np.zeros_like(self.qvals_2), - '|', color='red', markersize=25, markeredgewidth=2, label='qvals_2') - - # Plot already selected q-values from this session - if self.selected_qvals: - self.ax_qval.plot(self.selected_qvals, np.zeros_like(self.selected_qvals), - 'o', color='purple', markersize=10, label='Selected') - - # Redraw span patches - for patch in self.qval_spans: - self.ax_qval.add_patch(patch) - - if any([self.qvals_1 is not None, self.qvals_2 is not None, self.selected_qvals]): - self.ax_qval.legend() - - self.fig_qval.canvas.draw_idle() - - def _on_qval_span_select(self, xmin, xmax): - """Handle span selection for q-values.""" - # Select data within the span - mask = (self.data[:, 0] >= xmin) & (self.data[:, 0] <= xmax) - selected_data = self.data[mask, 0] - - if len(selected_data) > 0: - # Calculate median of selected range - median_qval = np.median(selected_data) - - # Add to selected q-values - self.selected_qvals.append(median_qval) - - # Create span patch for visualization - ylims = self.ax_qval.get_ylim() - span_patch = Rectangle((xmin, ylims[0]), xmax-xmin, ylims[1]-ylims[0], - alpha=0.2, color='purple') - self.qval_spans.append(span_patch) - - # Update plot - self._plot_qval_histogram(self.ax_qval.get_title()) - - print(f"Selected q-value: {median_qval:.6f} (from {len(selected_data)} points in range [{xmin:.6f}, {xmax:.6f}])") - - # Save to file - with open('selected_qvals.txt', 'a') as f: - f.write(f"{median_qval:.6f}\n") - - def _on_qval_clear_last(self, event): - """Remove the last selected q-value.""" - if self.selected_qvals: - removed = self.selected_qvals.pop() - if self.qval_spans: - self.qval_spans.pop() - print(f"Removed q-value: {removed:.6f}") - - # Update plot - self._plot_qval_histogram(self.ax_qval.get_title()) - - def _on_qval_key_press(self, event): - """Handle key press events for q-value selection.""" - if event.key == 'c': - # Clear last selection (same as button) - self._on_qval_clear_last(None) - elif event.key == 'd' or event.key == 'enter': - # Done selecting (same as button) - self._on_qval_done(None) - elif event.key == 'escape': - # Cancel and close without saving - self.selected_qvals = [] - plt.close(self.fig_qval) - - def _on_qval_done(self, event): - """Called when Done button is clicked for q-value selection.""" - print(f"Selected {len(self.selected_qvals)} q-values: {self.selected_qvals}") - plt.close(self.fig_qval) - - def select_triplets(self, title=None): - """Run the interactive selection process and return the selected triplets.""" - # Create the three windows with specific sizes - self.fig1, self.ax1 = plt.subplots(num='Step 1: Histogram Selection', figsize=(16,3)) - self.fig2, self.ax2 = plt.subplots(num='Step 2: 2D Selection', figsize=(16,3)) - self.fig3, (self.ax3a, self.ax3b, self.ax3c) = plt.subplots(1, 3, num='Step 3: Final Selection', - figsize=(16,3)) - - # Set window positions to stack them vertically - backend = plt.get_backend() - assert 'Tk' in backend - manager1 = self.fig1.canvas.manager - manager2 = self.fig2.canvas.manager - manager3 = self.fig3.canvas.manager - dpi = self.fig1.dpi - height1 = int(3 * dpi) # 3 inches * dpi - - # Position windows with some spacing - self.fig1.canvas.manager.window.wm_geometry("+100+50") - self.fig2.canvas.manager.window.wm_geometry(f"+100+{50 + height1 + 40}") - self.fig3.canvas.manager.window.wm_geometry(f"+100+{50 + 2*height1 + 80}") - - # Set fixed subplot sizes - #self.fig3.set_tight_layout(False) - self.fig1.subplots_adjust(bottom=.2) - self.fig2.subplots_adjust(bottom=.2) - self.fig3.subplots_adjust(bottom=.2) - - # Initialize the first window - self.show_histogram(title=title) - - # Set up the done button - done_ax = self.fig3.add_axes([0.9, 0.01, 0.09, 0.05]) - self.done_button = plt.Button(done_ax, 'Done') - self.done_button.on_clicked(self.on_done) - - # Flag to track when selection is complete - self.selection_done = False - - # Show plots and wait for user interaction - plt.show(block=True) - - # Close all figures - plt.close(self.fig1) - plt.close(self.fig2) - plt.close(self.fig3) - - # Return the selected triplets - return self.selected_triplets - - def compute_kde_maxima(self): - # Normalize data by bandwidths for anisotropic KDE - normalized_data = self.data / self.bandwidths[np.newaxis, :] - - # Create KernelDensity object - kde = KernelDensity(bandwidth=1, kernel='cosine') - print('fit') - kde.fit(normalized_data) - print('done fit') - - # Take a random subsample (5%) for evaluation - n_sample = max(int(len(normalized_data) * 0.05), 100000) # At least 1000 points - n_sample = 200000 - n_sample = min(n_sample, len(normalized_data)) # Can't sample more than we have - print(f'{n_sample=}') - - # Random sampling without replacement - sample_indices = np.random.choice(len(normalized_data), size=n_sample, replace=False) - sample_data = normalized_data[sample_indices] - - # Evaluate KDE at sampled points - print('score') - sample_densities = np.exp(kde.score_samples(sample_data)) - print('done score') - - # Sort sampled points by density value - sorted_indices = np.argsort(sample_densities)[::-1] - - # Filter to avoid maxima that are too close together - safety_factor = 2.0 # How many bandwidths apart maxima should be - min_distances = self.bandwidths * safety_factor - - # Initialize list to store maxima - maxima_indices = [] - maxima_points = [] - maxima_densities = [] - - # Function to be minimized (negative density) - def negative_density(x): - return -np.exp(kde.score_samples([x])[0]) - - # Process points in order of decreasing density - for i, idx in enumerate(sorted_indices): - if i%1000==0: - print('processed', i, ', kept', len(maxima_points)) - if len(maxima_points) >= self.n_maxima: - break - - initial_point = sample_data[idx] - q1, q2, th = initial_point - if np.abs(q1-q2) < 2: continue - if th<8 or th>160: continue - - optimized_point = initial_point - optimized_density = 1 -# # Run optimization to find true maximum -# result = minimize( -# negative_density, -# initial_point, -# method='BFGS', -# #options={'gtol': 1e-5} # Gradient tolerance for convergence -# ) -# -# if result.success: -# optimized_point = result.x -# optimized_density = -result.fun # Negate back to get positive density - - - # Check if this point is far enough from all accepted maxima - too_close = False - for accepted_point in maxima_points: - dist = np.linalg.norm(optimized_point - accepted_point) - if dist < 5: - too_close = True - break - - sus = False - q1,q2,th = optimized_point - if np.abs(q1-q2) < 2: sus = True - if th < 5 or th > 160: sus = True - - if not too_close and not sus: - maxima_points.append(optimized_point) - maxima_densities.append(optimized_density) - - # Store the maxima and their density values - self.kde_maxima = np.array(maxima_points) * self.bandwidths[np.newaxis, :] - self.kde_values = np.array(maxima_densities) - with open('autopeaks.txt', 'w') as f: - for x in self.kde_maxima: - print(round(x[0], 4), round(x[1], 4), round(x[2], 2), file=f) - - - print(f"Found {len(self.kde_maxima)} KDE maxima from {n_sample} sampled points") - - def compute_pf_maxima(self): - # Define bins with 2x oversampling - bin_width = self.bandwidths / 2.0 - - bins = [ - np.arange(0.05, 0.5 + bin_width[0], bin_width[0]), - np.arange(0.05, 0.5 + bin_width[1], bin_width[1]), - np.arange(10, 160 + bin_width[2], bin_width[2]) - ] - - print(f'Creating histogram with shape: {[len(b)-1 for b in bins]}') - - from scipy.spatial import cKDTree - - # Build tree once at start - tree = cKDTree(self.data / self.bandwidths) - - # Create 3D histogram - hist, edges = np.histogramdd(self.data, bins=bins) - - # Smooth with gaussian (sigma ~ 1 bin to merge nearby peaks) - from scipy import ndimage - hist_smooth = ndimage.gaussian_filter(hist, sigma=1) - - # Find local maxima - max_filtered = ndimage.maximum_filter(hist_smooth, size=10) - peaks = (hist_smooth == max_filtered) - - # Threshold to remove noise peaks - threshold = hist_smooth.max() * 0.01 # Adjust as needed - peaks &= (hist_smooth > threshold) - - print(f'Found {peaks.sum()} initial peaks') - - # Get peak coordinates in bin indices - peak_indices = np.argwhere(peaks) - peak_values = hist_smooth[peaks] - - # Convert bin indices to actual coordinates - peak_coords = np.array([ - edges[0][peak_indices[:, 0]] + bin_width[0]/2, - edges[1][peak_indices[:, 1]] + bin_width[1]/2, - edges[2][peak_indices[:, 2]] + bin_width[2]/2 - ]).T - - # Apply your domain filters - valid_mask = np.ones(len(peak_coords), dtype=bool) - - q1, q2, th = peak_coords[:, 0], peak_coords[:, 1], peak_coords[:, 2] - valid_mask &= (np.abs(q1 - q2) >= 2 * self.bandwidths[0]) - valid_mask &= (th >= 8) & (th <= 160) - - peak_coords = peak_coords[valid_mask] - peak_values = peak_values[valid_mask] - sortkey = peak_values - #sortkey = -1* (peak_coords[:,0] + peak_coords[:,1]) - - # Sort by density - sorted_indices = np.argsort(sortkey)[::-1] - - # Filter by minimum distance (your safety_factor logic) - min_dist = 10.0 # In your normalized space this was 5 - final_peaks = [] - final_values = [] - - for idx in sorted_indices: - if len(final_peaks) >= self.n_maxima: - break - - candidate = peak_coords[idx] - - # Check if it's a true local maximum - if len(final_peaks) > 0: - # Option 1: Saddle test - is_separate = all( - has_saddle_between( - candidate, accepted, hist_smooth, edges, bin_width, prominence_threshold=0.8 - ) - for accepted in final_peaks - if np.linalg.norm((candidate - accepted) / bin_width) < 180 # Normalized distance - ) - - # Option 2: Gradient test (faster) - # is_separate = is_local_maximum(candidate, hist_smooth, edges, bin_width, final_peaks) - - if not is_separate: - continue - - final_peaks.append(candidate) - final_values.append(peak_values[idx]) - - # Refine peak locations by fitting Gaussians - refined_peaks = [] - - for i, peak in enumerate(final_peaks): - if i % 100 == 0: - print(f'Refining peak {i}/{len(final_peaks)}') - - # Select points within 2x bandwidth - normalized_peak = peak / self.bandwidths - indices = tree.query_ball_point(normalized_peak, r=2.0) - nearby_points = self.data[indices] - - if len(nearby_points) >= 5: # Need enough points to fit - refined_peak = fit_gaussian_peak( - nearby_points, peak, self.bandwidths - ) - refined_peaks.append(refined_peak) - else: - print('fallback') - refined_peaks.append(peak) # Not enough points, keep grid location - - self.kde_maxima = np.array(refined_peaks) - self.kde_values = np.array(final_values) - - # Final sorting - sort_vals = self.kde_maxima[:, 0] + self.kde_maxima[:, 1] - final_indices = np.argsort(sort_vals)[:self.n_shortest] - self.kde_maxima = self.kde_maxima[final_indices] - self.kde_values = self.kde_values[final_indices] - - - def compute_slice_maxima(self, n_q1_peaks=30, min_q1=0.10, max_q1=0.35, - slice_width=0.002, max_peaks_per_slice=100, - max_peaks=500, verbose=True): - """ - Find cluster centers using slice-based 2D KDE approach. - - This method finds q1 peaks using 1D KDE, then performs 2D KDE - peak finding within each q1 slice, followed by 3D refinement. - - Final peaks are selected by score / sqrt(q1 * q2) to prioritize - low-q peaks which are most valuable for unit cell determination. - - Parameters - ---------- - n_q1_peaks : int - Maximum number of q1 slices to process - min_q1, max_q1 : float - Range for q1 peak detection - slice_width : float - Width of each q1 slice - max_peaks_per_slice : int - Maximum peaks to find per slice - verbose : bool - Print progress - """ - from scipy.ndimage import maximum_filter - from scipy.optimize import minimize_scalar - from scipy.spatial import cKDTree - - data = self.data - bandwidths_3d = self.bandwidths - bandwidths_2d = np.array([bandwidths_3d[1], bandwidths_3d[2]]) - - # Step 1: Find q1 peaks using 1D Gaussian KDE - if verbose: - print("Finding q1 peaks...") - - q1_data = data[:, 0] - q1_mask = (q1_data >= min_q1) & (q1_data <= max_q1) - q1_filtered = q1_data[q1_mask] - - # Subsample for KDE - np.random.seed(42) - n_sample = min(1000000, len(q1_filtered)) - sample_idx = np.random.choice(len(q1_filtered), n_sample, replace=False) - q1_sample = q1_filtered[sample_idx].reshape(-1, 1) - - # Fit 1D KDE - kde_1d = KernelDensity(bandwidth=0.0002, kernel='cosine') - kde_1d.fit(q1_sample) - - # Find peaks on fine grid - n_eval = 10000 - q1_grid = np.linspace(min_q1, max_q1, n_eval) - density_1d = np.exp(kde_1d.score_samples(q1_grid.reshape(-1, 1))) - - footprint = np.ones(5) - local_max = maximum_filter(density_1d, footprint=footprint) - peaks_mask = (density_1d == local_max) & (density_1d > np.percentile(density_1d, 50)) - peak_q1s_raw = q1_grid[np.where(peaks_mask)[0]] - - # Refine peak positions - def neg_density(q1): - return -kde_1d.score_samples([[q1]])[0] - - refined_q1s = [] - for q1 in peak_q1s_raw: - result = minimize_scalar(neg_density, bounds=(q1-0.0005, q1+0.0005), method='bounded') - refined_q1s.append(result.x) - refined_q1s = np.array(refined_q1s) - - # Score and select with spacing - peak_densities = np.exp(kde_1d.score_samples(refined_q1s.reshape(-1, 1))) - scores = peak_densities / (refined_q1s ** 0.5) - sort_idx = np.argsort(scores)[::-1] - q1s_sorted = refined_q1s[sort_idx] - - min_spacing = 0.002 - q1_centers = [] - for q1 in q1s_sorted: - if len(q1_centers) >= n_q1_peaks: - break - if all(abs(q1 - s) >= min_spacing for s in q1_centers): - q1_centers.append(q1) - q1_centers = np.sort(q1_centers) - - if verbose: - print(f"Found {len(q1_centers)} q1 peaks") - - # Step 2: Process each q1 slice - all_peaks = [] - all_scores = [] - - for i, q1_center in enumerate(q1_centers): - q1_min = q1_center - slice_width / 2 - q1_max = q1_center + slice_width / 2 - - slice_mask = (data[:, 0] >= q1_min) & (data[:, 0] <= q1_max) - slice_data = data[slice_mask] - - if len(slice_data) < 50: - continue - - if verbose: - print(f"Slice {i+1}/{len(q1_centers)}: q1=[{q1_min:.4f}, {q1_max:.4f}], {len(slice_data)} points", end="") - - # Filter diagonal and theta - diag_threshold = 2.0 - theta_range = (8, 160) - q_diff_norm = np.abs(slice_data[:, 0] - slice_data[:, 1]) / bandwidths_3d[0] - off_diag_mask = q_diff_norm >= diag_threshold - theta_mask = (slice_data[:, 2] >= theta_range[0]) & (slice_data[:, 2] <= theta_range[1]) - slice_data_filtered = slice_data[off_diag_mask & theta_mask] - - if len(slice_data_filtered) < 50: - if verbose: - print(" -> 0 peaks (filtered)") - continue - - # 3D KDE and tree on full slice - normalized_3d = slice_data / bandwidths_3d - kde_3d = KernelDensity(bandwidth=1, kernel='cosine') - kde_3d.fit(normalized_3d) - tree_3d = cKDTree(normalized_3d) - - # 2D KDE on filtered subsample - max_sample_2d = 30000 - if len(slice_data_filtered) > max_sample_2d: - sample_idx_2d = np.random.choice(len(slice_data_filtered), max_sample_2d, replace=False) - slice_2d_sample = slice_data_filtered[sample_idx_2d, 1:3] - else: - slice_2d_sample = slice_data_filtered[:, 1:3] - - normalized_2d_sample = slice_2d_sample / bandwidths_2d - kde_2d = KernelDensity(bandwidth=1, kernel='cosine') - kde_2d.fit(normalized_2d_sample) - scores_2d = kde_2d.score_samples(normalized_2d_sample) - tree_2d = cKDTree(normalized_2d_sample) - - # Find 2D peaks - def refine_2d(point_norm): - current = point_norm.copy() - for _ in range(5): - idx = tree_2d.query_ball_point(current, r=1.5) - if len(idx) < 5: - break - nearby = normalized_2d_sample[idx] - nearby_scores = scores_2d[idx] - weights = np.exp(nearby_scores - nearby_scores.max()) - new_pos = np.average(nearby, axis=0, weights=weights) - if np.linalg.norm(new_pos - current) < 0.01: - break - current = new_pos - return current - - top_k = min(5000, len(slice_2d_sample)) - top_idx = np.argsort(scores_2d)[-top_k:] - - min_distance_2d = 3.0 - peaks_2d_norm = [] - for idx in np.argsort(scores_2d[top_idx])[::-1]: - point_idx = top_idx[idx] - refined = refine_2d(normalized_2d_sample[point_idx]) - is_close = any(np.linalg.norm(refined - ex) < min_distance_2d for ex in peaks_2d_norm) - if not is_close: - peaks_2d_norm.append(refined) - if len(peaks_2d_norm) >= max_peaks_per_slice: - break - - # Refine in 3D - def refine_3d(point_3d, max_neighbors=300): - current = point_3d.copy() / bandwidths_3d - for _ in range(10): - idx = tree_3d.query_ball_point(current, r=2.0) - if len(idx) < 5: - break - idx = np.array(idx) - if len(idx) > max_neighbors: - idx = idx[np.random.choice(len(idx), max_neighbors, replace=False)] - nearby = normalized_3d[idx] - nearby_scores = kde_3d.score_samples(nearby) - weights = np.exp(nearby_scores - nearby_scores.max()) - new_pos = np.average(nearby, axis=0, weights=weights) - if np.linalg.norm(new_pos - current) < 0.01: - break - current = new_pos - return current * bandwidths_3d - - min_distance_3d = 3.0 - slice_peaks = [] - slice_scores = [] - - for p2d_norm in peaks_2d_norm: - p2d = p2d_norm * bandwidths_2d - dq2 = np.abs(slice_data[:, 1] - p2d[0]) - dtheta = np.abs(slice_data[:, 2] - p2d[1]) - nearby_mask = (dq2 < 0.003) & (dtheta < 4.0) - - if nearby_mask.sum() < 5: - continue - - nearby_scores_3d = kde_3d.score_samples(normalized_3d[nearby_mask]) - best_nearby = slice_data[nearby_mask][np.argmax(nearby_scores_3d)] - - refined_3d = refine_3d(best_nearby) - refined_norm = refined_3d / bandwidths_3d - - # Check theta range - if refined_3d[2] < theta_range[0] or refined_3d[2] > theta_range[1]: - continue - - # Check distance to existing peaks - is_close = any(np.linalg.norm(refined_norm - (ex / bandwidths_3d)) < min_distance_3d - for ex in slice_peaks) - if is_close: - continue - - refined_score = kde_3d.score_samples(refined_norm.reshape(1, -1))[0] - slice_peaks.append(refined_3d) - slice_scores.append(refined_score) - - all_peaks.extend(slice_peaks) - all_scores.extend(slice_scores) - - if verbose: - print(f" -> {len(slice_peaks)} peaks") - - # Deduplicate across slices - if len(all_peaks) > 0: - all_peaks = np.array(all_peaks) - all_scores = np.array(all_scores) - - sort_idx = np.argsort(all_scores)[::-1] - all_peaks = all_peaks[sort_idx] - all_scores = all_scores[sort_idx] - - unique_peaks = [] - unique_scores = [] - q_tol, theta_tol = 0.0005, 2.0 - - for p, s in zip(all_peaks, all_scores): - is_dup = any( - abs(p[0] - u[0]) < q_tol and - abs(p[1] - u[1]) < q_tol and - abs(p[2] - u[2]) < theta_tol - for u in unique_peaks - ) - if not is_dup: - unique_peaks.append(p) - unique_scores.append(s) - - unique_peaks = np.array(unique_peaks) - unique_scores = np.array(unique_scores) - - # Select top peaks prioritizing low q1*q2 (log-space adjustment) - if max_peaks is not None and len(unique_peaks) > max_peaks: - q1_vals = unique_peaks[:, 0] - q2_vals = unique_peaks[:, 1] - # score - 0.5*log(q1*q2) is the log-space equivalent of score/sqrt(q1*q2) - selection_scores = unique_scores - 0.5 * np.log(q1_vals * q2_vals) - top_idx = np.argsort(selection_scores)[-max_peaks:] - unique_peaks = unique_peaks[top_idx] - unique_scores = unique_scores[top_idx] - - if verbose: - print(f"Selected top {max_peaks} peaks (prioritizing low-q)") - - self.kde_maxima = unique_peaks - self.kde_values = unique_scores - else: - self.kde_maxima = np.array([]).reshape(0, 3) - self.kde_values = np.array([]) - - if verbose: - print(f"Final: {len(self.kde_maxima)} maxima") - - def show_histogram(self, title=None): - self.ax1.clear() - plotrange=self.qmin, self.qmax - self.ax1.set_xlim(plotrange) - hist, edges = np.histogram(self.data[:,0], bins=2000, range=plotrange) - self.ax1.plot(edges[:-1], hist) - if title is None: - title = "Select range in histogram" - self.ax1.set_title(title) - self.ax1.set_xlabel('q1') - self.ax1.set_ylabel('counts') - - # Add tick marks for KDE maxima if available - if hasattr(self, 'kde_maxima'): - self.ax1.plot(self.kde_maxima[:, 0], np.zeros_like(self.kde_maxima[:, 0]), - '|', color='green', markersize=20) - - # Add tick marks for specified q-values - if self.qvals_1 is not None: - self.ax1.plot(self.qvals_1, np.zeros_like(self.qvals_1), - '|', color='blue', markersize=25, markeredgewidth=2) - if self.qvals_2 is not None: - self.ax1.plot(self.qvals_2, np.zeros_like(self.qvals_2), - '|', color='red', markersize=25, markeredgewidth=2) - - if self.sb1_qvals is not None: - self.ax1.plot(self.sb1_qvals, np.zeros_like(self.sb1_qvals), - '|', color='orange', markersize=25, markeredgewidth=2) - - - self.span = SpanSelector( - self.ax1, - self.on_span_select, - 'horizontal', - useblit=True, - props=dict(alpha=0.5, facecolor='red') - ) - - # Redraw the figure - self.fig1.canvas.draw_idle() - - def on_span_select(self, xmin, xmax): - # Remove previous span patch if it exists - if self.span_patch is not None: - try: - self.span_patch.remove() - except NotImplementedError: - pass - self.span_patch = None - - # Create new span patch - ylims = self.ax1.get_ylim() - self.span_patch = Rectangle((xmin, ylims[0]), xmax-xmin, ylims[1]-ylims[0], - alpha=0.2, color='red') - self.ax1.add_patch(self.span_patch) - - # Select data within the span - mask = (self.data[:, 0] >= xmin) & (self.data[:, 0] <= xmax) - self.current_selection = self.data[mask] - - # Select KDE maxima within the span - if hasattr(self, 'kde_maxima'): - kde_mask = (self.kde_maxima[:, 0] >= xmin) & (self.kde_maxima[:, 0] <= xmax) - self.current_kde_selection = self.kde_maxima[kde_mask] - - # Select predicted points in the span - if self.points is not None: - points_mask = (self.points[:,0] >= xmin) & (self.points[:,0] <= xmax) - self.current_points_selection = self.points[points_mask] - else: - self.current_points_selection = None - - # Show the second window - self.show_scatter_2d() - - # Redraw both figures - self.fig1.canvas.draw_idle() - self.fig2.canvas.draw_idle() - - def show_scatter_2d(self): - self.ax2.clear() - plotrange = self.qmin, self.qmax - self.ax2.set_xlim(plotrange) - self.ax2.set_xlabel('q2') - self.ax2.set_ylabel('theta') - if self.current_selection is not None and len(self.current_selection) > 0: - self.ax2.scatter( - self.current_selection[:, 1], - self.current_selection[:, 2], - s=2, alpha=.2) - self.ax2.set_title(f'Draw box to select points (n={len(self.current_selection)})') - - # Plot KDE maxima if available - if hasattr(self, 'current_kde_selection') and len(self.current_kde_selection) > 0: - self.ax2.scatter(self.current_kde_selection[:, 1], self.current_kde_selection[:, 2], - color='red', marker='x', s=50, label='KDE maxima') - - - # Add vertical lines for q-values (second dimension) - if self.qvals_2 is not None: - ylim = self.ax2.get_ylim() - for q in self.qvals_2: - if q >= self.current_selection[:, 1].min() and q <= self.current_selection[:, 1].max(): - self.ax2.axvline(q, color='blue', linestyle='--', alpha=0.5) - if self.qvals_1 is not None and self.points is None: - for q in self.qvals_1: - if q >= self.current_selection[:, 1].min() and q <= self.current_selection[:, 1].max(): - self.ax2.axvline(q, color='red', linestyle='--', alpha=0.5) - if self.current_points_selection is not None: - self.ax2.scatter(self.current_points_selection[:,1], self.current_points_selection[:,2], - color='orange', marker='o', s=20, label='points') - - # Create RectangleSelector only if it doesn't exist already - if not hasattr(self, 'rect') or self.rect is None: - self.rect = RectangleSelector( - self.ax2, - self.on_rect_select, - useblit=True, - props=dict(facecolor='red', alpha=0.2) - ) - - # If there was a previous rectangle, redraw it - if self.rect_patch is not None: - self.ax2.add_patch(self.rect_patch) - else: - self.ax2.set_title('No points selected') - - self.fig2.canvas.mpl_connect('key_press_event', self.on_key_press) - self.fig2.canvas.draw_idle() - - def on_rect_select(self, eclick, erelease): - x1, y1 = eclick.xdata, eclick.ydata - x2, y2 = erelease.xdata, erelease.ydata - - # Remove previous rectangle if it exists - if self.rect_patch is not None: - self.rect_patch.remove() - - # Create new rectangle patch - self.rect_patch = Rectangle((min(x1, x2), min(y1, y2)), - abs(x2-x1), abs(y2-y1), - alpha=0.2, color='red') - self.ax2.add_patch(self.rect_patch) - self.fig2.canvas.draw_idle() - - # Select data within the rectangle - mask = ( - (self.current_selection[:, 1] >= min(x1, x2)) & - (self.current_selection[:, 1] <= max(x1, x2)) & - (self.current_selection[:, 2] >= min(y1, y2)) & - (self.current_selection[:, 2] <= max(y1, y2)) - ) - self.final_selection = self.current_selection[mask] - - # Select KDE maxima within the rectangle - if hasattr(self, 'current_kde_selection'): - kde_mask = ( - (self.current_kde_selection[:, 1] >= min(x1, x2)) & - (self.current_kde_selection[:, 1] <= max(x1, x2)) & - (self.current_kde_selection[:, 2] >= min(y1, y2)) & - (self.current_kde_selection[:, 2] <= max(y1, y2)) - ) - self.final_kde_selection = self.current_kde_selection[kde_mask] - - # Show the third window - self.show_final_plots() - - def show_final_plots(self): - self.update_means() - - # Clear all axes - for ax in [self.ax3a, self.ax3b, self.ax3c]: - ax.clear() - #ax.set_aspect('equal', adjustable='datalim') - - # Find the limits that contain all selected points - x_min, x_max = self.final_selection[:, 0].min(), self.final_selection[:, 0].max() - y_min, y_max = self.final_selection[:, 1].min(), self.final_selection[:, 1].max() - z_min, z_max = self.final_selection[:, 2].min(), self.final_selection[:, 2].max() - - # Add a small margin - margin = 0.1 - x_range = x_max - x_min - y_range = y_max - y_min - z_range = z_max - z_min - -# x_min -= margin * x_range -# x_max += margin * x_range -# y_min -= margin * y_range -# y_max += margin * y_range -# z_min -= margin * z_range -# z_max += margin * z_range - - # Set the limits for each plot - self.ax3a.set_xlim(x_min, x_max) - self.ax3a.set_ylim(y_min, y_max) - self.ax3a.set_xlabel('q1') - self.ax3a.set_ylabel('q2') - self.ax3b.set_xlabel('q1') - self.ax3b.set_ylabel('theta') - self.ax3c.set_xlabel('q2') - self.ax3c.set_ylabel('theta') - - self.ax3b.set_xlim(x_min, x_max) - self.ax3b.set_ylim(z_min, z_max) - - self.ax3c.set_xlim(y_min, y_max) - self.ax3c.set_ylim(z_min, z_max) - - # Plot all points within the final selection - self.ax3a.scatter(self.final_selection[:, 0], self.final_selection[:, 1], color='tab:blue', - s=5, alpha=.3) - self.ax3b.scatter(self.final_selection[:, 0], self.final_selection[:, 2], color='tab:blue', - s=5, alpha=.3) - self.ax3c.scatter(self.final_selection[:, 1], self.final_selection[:, 2], color='tab:blue', - s=5, alpha=.3) - - # Plot KDE maxima - if hasattr(self, 'final_kde_selection') and len(self.final_kde_selection) > 0: - self.ax3a.scatter(self.final_kde_selection[:, 0], self.final_kde_selection[:, 1], - color='red', marker='x', s=50) - self.ax3b.scatter(self.final_kde_selection[:, 0], self.final_kde_selection[:, 2], - color='red', marker='x', s=50) - self.ax3c.scatter(self.final_kde_selection[:, 1], self.final_kde_selection[:, 2], - color='red', marker='x', s=50) - - # Plot means - self.ax3a.plot(self.means[0], self.means[1], 'ro', markersize=6) - self.ax3b.plot(self.means[0], self.means[2], 'ro', markersize=6) - self.ax3c.plot(self.means[1], self.means[2], 'ro', markersize=6) - - # Add rectangle selectors to all plots - self.rect_final = [ - RectangleSelector(ax, self.on_final_rect_select, useblit=True, - props=dict(facecolor='red', alpha=0.2)) - for ax in [self.ax3a, self.ax3b, self.ax3c] - ] - - self.fig3.canvas.mpl_connect('key_press_event', self.on_key_press) - self.fig3.canvas.draw_idle() - - - - def on_final_rect_select(self, eclick, erelease): - x1, y1 = eclick.xdata, eclick.ydata - x2, y2 = erelease.xdata, erelease.ydata - - # Get the current axis - ax = eclick.inaxes - - # Create mask based on which plot was clicked - if ax == self.ax3a: - mask = ( - (self.final_selection[:, 0] >= min(x1, x2)) & - (self.final_selection[:, 0] <= max(x1, x2)) & - (self.final_selection[:, 1] >= min(y1, y2)) & - (self.final_selection[:, 1] <= max(y1, y2)) - ) - elif ax == self.ax3b: - mask = ( - (self.final_selection[:, 0] >= min(x1, x2)) & - (self.final_selection[:, 0] <= max(x1, x2)) & - (self.final_selection[:, 2] >= min(y1, y2)) & - (self.final_selection[:, 2] <= max(y1, y2)) - ) - elif ax == self.ax3c: - mask = ( - (self.final_selection[:, 1] >= min(x1, x2)) & - (self.final_selection[:, 1] <= max(x1, x2)) & - (self.final_selection[:, 2] >= min(y1, y2)) & - (self.final_selection[:, 2] <= max(y1, y2)) - ) - - # Save the limits - axl = self.ax3a.get_xlim() - ayl = self.ax3a.get_ylim() - bxl = self.ax3b.get_xlim() - byl = self.ax3b.get_ylim() - cxl = self.ax3c.get_xlim() - cyl = self.ax3c.get_ylim() - # Clear all plots including mean markers - for ax in [self.ax3a, self.ax3b, self.ax3c]: - ax.clear() - #ax.set_aspect('equal', adjustable='datalim') - # Restore the original limits - if ax == self.ax3a: - ax.set_xlim(axl) - ax.set_ylim(ayl) - elif ax == self.ax3b: - ax.set_xlim(bxl) - ax.set_ylim(byl) - else: - ax.set_xlim(cxl) - ax.set_ylim(cyl) - - # Update means based on the new selection - selected_points = self.final_selection[mask] - self.means = np.mean(selected_points, axis=0) - - # Replot everything - self.ax3a.scatter(self.final_selection[~mask, 0], self.final_selection[~mask, 1], - color='gray', s=5, alpha=0.3) - self.ax3a.scatter(self.final_selection[mask, 0], self.final_selection[mask, 1], - s=5, alpha=.3, color='tab:blue') - self.ax3a.plot(self.means[0], self.means[1], 'ro', markersize=6) - - self.ax3b.scatter(self.final_selection[~mask, 0], self.final_selection[~mask, 2], - color='gray', s=5, alpha=0.3) - self.ax3b.scatter(self.final_selection[mask, 0], self.final_selection[mask, 2], - s=5, alpha=.3, color='tab:blue') - self.ax3b.plot(self.means[0], self.means[2], 'ro', markersize=6) - - self.ax3c.scatter(self.final_selection[~mask, 1], self.final_selection[~mask, 2], - color='gray', s=5, alpha=0.3) - self.ax3c.scatter(self.final_selection[mask, 1], self.final_selection[mask, 2], - s=5, alpha=.3, color='tab:blue') - self.ax3c.plot(self.means[1], self.means[2], 'ro', markersize=6) - - # Plot KDE maxima - if hasattr(self, 'final_kde_selection') and len(self.final_kde_selection) > 0: - self.ax3a.scatter(self.final_kde_selection[:, 0], self.final_kde_selection[:, 1], - color='red', marker='x', s=50) - self.ax3b.scatter(self.final_kde_selection[:, 0], self.final_kde_selection[:, 2], - color='red', marker='x', s=50) - self.ax3c.scatter(self.final_kde_selection[:, 1], self.final_kde_selection[:, 2], - color='red', marker='x', s=50) - - # Plot means - self.ax3a.plot(self.means[0], self.means[1], 'ro', markersize=6) - self.ax3b.plot(self.means[0], self.means[2], 'ro', markersize=6) - self.ax3c.plot(self.means[1], self.means[2], 'ro', markersize=6) - - self.ax3a.set_xlabel('q1') - self.ax3a.set_ylabel('q2') - self.ax3b.set_xlabel('q1') - self.ax3b.set_ylabel('theta') - self.ax3c.set_xlabel('q2') - self.ax3c.set_ylabel('theta') - - self.fig3.canvas.draw_idle() - - def update_means(self): - self.means = np.mean(self.final_selection, axis=0) - - def on_key_press(self, event): - if event.key == 'a': - # Add mean to selected triplets - triplet = [self.means[0], self.means[1], self.means[2]] - self.selected_triplets.append(triplet) - - print(f"Selected point {self.means[0]:.6f} {self.means[1]:.6f} {self.means[2]:.6f}. " - f"Total triplets: {len(self.selected_triplets)}") - - # Also append to file if desired - with open('cluster_means.txt', 'a') as f: - np.savetxt(f, [self.means], fmt='%.6f') - - if self.sb1_callback is not None: - self.sb1_qvals, self.points = self.sb1_callback(triplet, self.recon) - self.show_histogram(title=self.ax1.get_title()) - - def on_done(self, event): - """Called when the Done button is clicked.""" - self.selection_done = True - plt.close('all') # Close all open figures - import gc;gc.collect() - - -# Example usage: -if __name__ == "__main__": - # Janky format - alldata = [] - for f in sys.argv[1:]: - data = np.load(f)['triplets'][:,1:4] - data[:,0] = 1/data[:,0] - data[:,1] = 1/data[:,1] - data2 = np.vstack((data[:,1], data[:,0], data[:,2])).transpose() - data = np.vstack((data, data2)) - alldata.append(data) - alldata = np.vstack(alldata) - #alldata = alldata[:100000] - - bandwidths = [0.001, 0.001, 1.0] - - clusterer = ManualClusterer(alldata, bandwidths=bandwidths, n_maxima=500) - clusterer.select_triplets() - diff --git a/xfel/small_cell/command_line/index2.py b/xfel/small_cell/command_line/index2.py deleted file mode 100644 index ed572e79c4f..00000000000 --- a/xfel/small_cell/command_line/index2.py +++ /dev/null @@ -1,2640 +0,0 @@ -import sys -import numpy as np -from dataclasses import dataclass -from typing import List, Tuple, Optional -import copy -import matplotlib.pyplot as plt -from matplotlib.colors import LinearSegmentedColormap -from tqdm import tqdm -from numba import njit -from scipy.optimize import minimize_scalar -from scipy.optimize import minimize -import itertools -from cluster2 import ManualClusterer -from cctbx import uctbx, crystal -from cctbx.sgtbx.lattice_symmetry import metric_subgroups - -class SpotPair: - - def __init__(self, q1: float, q2: float, theta: float, preserve_order=False): - """Just q1, q2, theta. Theta is in radians.""" - # Ensure q1 <= q2 - if not preserve_order and q1 > q2: - q1, q2 = q2, q1 - self.q1 = q1 - self.q2 = q2 - self.theta = theta - - def area(self) -> float: - """Calculate twice the area of the triangle formed by the vectors.""" - v1 = np.array([self.q1, 0.0]) - v2 = self.q2 * np.array([np.cos(self.theta), np.sin(self.theta)]) - return abs(np.cross(v1, v2)) - - - -class VectorPairMatch: - @classmethod - def from_pairs(cls, gen_pair, obs_pair): - return - -class PairMatch2d(VectorPairMatch): - def __init__(self, hkl1, q1, hkl2, q2, theta_rad): - self.hkl1 = hkl1 - self.hkl2 = hkl2 - self.q1 = q1 - self.q2 = q2 - self.theta_rad = theta_rad - self.is_outlier = False - -class OneVectorMatch(VectorPairMatch): - """ - q1 is the indexed vector and hkl1 is the corresponding indices in the sublattice. - """ - def __init__(self, hkl1, q1, q2, theta_rad): - self.hkl1 = hkl1 - self.q1 = q1 - self.q2 = q2 - self.theta_rad = theta_rad - - -# A couple helper functions - -def angle_between(v1, v2): - v1_u = v1/np.linalg.norm(v1) - v2_u = v2/np.linalg.norm(v2) - return np.degrees(np.arccos(np.clip(np.dot(v1_u, v2_u), -1.0, 1.0))) - -def third_vector(q, v1, v2, theta1, theta2): - """ - Compute a vector v3 with length q that forms angles theta1 and theta2 - with vectors v1 and v2 respectively. - - Parameters: - q (float): Desired length of the output vector - v1 (array-like): First reference vector - v2 (array-like): Second reference vector - theta1 (float): Desired angle with v1 (in radians) - theta2 (float): Desired angle with v2 (in radians) - - Returns: - numpy.ndarray: The computed vector v3 - """ - # Convert inputs to numpy arrays and normalize vectors - v1 = np.array(v1, dtype=float) - v2 = np.array(v2, dtype=float) - v1_norm = np.linalg.norm(v1) - v2_norm = np.linalg.norm(v2) - v1 = v1 / v1_norm - v2 = v2 / v2_norm - - # The vector we're looking for can be written as a linear combination - # of v1, v2, and their cross product: v3 = a*v1 + b*v2 + c*(v1 × v2) - - # First, get the cross product and normalize it - v1xv2 = np.cross(v1, v2) - if np.allclose(v1xv2, 0): - raise ValueError("Input vectors are parallel, solution is not unique") - v1xv2 = v1xv2 / np.linalg.norm(v1xv2) - - # The conditions are: - # q*cos(theta1) = a + b*cos(gamma) - # q*cos(theta2) = a*cos(gamma) + b - # a^2 + b^2 + c^2 = q^2 - # where gamma is the angle between v1 and v2 - - cos_gamma = np.dot(v1, v2) - sin_gamma = np.sqrt(1 - cos_gamma**2) - - # Solve for a and b - A = np.array([[1, cos_gamma], - [cos_gamma, 1]]) - b = q * np.array([np.cos(theta1), np.cos(theta2)]) - - try: - a, b = np.linalg.solve(A, b) - - # Now solve for c using the Pythagorean theorem - c_sq = q**2 - (a**2 + b**2 + 2*a*b*cos_gamma) - if c_sq < 0: - raise ValueError("No solution exists for these angles") - c = np.sqrt(c_sq) - - # There are two possible solutions (±c) - # We'll return the positive c solution - v3 = a*v1 + b*v2 + c*v1xv2 - - return v3 - - except np.linalg.LinAlgError: - raise ValueError("No solution exists for these angles") - - -def find_best_third_vector(v3_candidates: Tuple[np.ndarray, np.ndarray], - sub_basis: np.ndarray) -> np.ndarray: - """From two possible v3 positions, generate all lattice-equivalent points - and choose the closest to origin. - - Args: - v3_candidates: Two possible positions for v3 - sub_basis: Current 2x2 sublattice basis - - Returns: - The best choice for the third basis vector - """ - a, b = (np.hstack((x,0)) for x in sub_basis) - best_v3 = None - min_length = float('inf') - - for v3 in v3_candidates: - # Generate all equivalent points v3 + ha + kb - # Check a generous range of h,k values - for h in range(-3, 4): - for k in range(-3, 4): - v3_equiv = v3 + h*a + k*b - length = np.linalg.norm(v3_equiv) - #print(h, k, round(length,5)) - - if length < min_length: - min_length = length - best_v3 = v3_equiv - - return best_v3 - - -class Basis: - def __init__(self, vectors: np.ndarray, qmax: float, - q_tolerance: float = 0.001, theta_tol_degrees: float = 1.0, - symmetry=None, centering=None): - self.vectors = vectors - self.qmax = qmax - self.q_tolerance = q_tolerance - self.theta_tolerance = np.radians(theta_tol_degrees) - self.points = None - self.point_indices = None - self.pairs = None - self.indexed_pairs = [] - self.symmetry = symmetry - self.centering = centering or 'P' - - @classmethod - def from_vectors(cls, vectors: np.ndarray, reduce=True, qmax:float = 0.5, q_tol: float = 0.001, - theta_tol_deg: float = 1.0, symmetry=None, centering=None): - assert vectors.shape in ((2,2), (3,3)) - if vectors.shape == (2,2): - temp_basis = Basis2d(vectors, qmax, q_tol, theta_tol_deg) - - if not reduce: # Short circuit - return temp_basis - - temp_basis.generate_points_and_pairs_fast() - points = temp_basis.points - - reduced_vectors = [] - distances = np.linalg.norm(points, axis=1) - sort_idx = np.argsort(distances) - for i in sort_idx: - p = points[i] - if len(reduced_vectors) == 0 and np.linalg.norm(p) > 1e-6: - reduced_vectors.append(p) - elif len(reduced_vectors) == 1: - cross_prod = np.cross(reduced_vectors[0], p) - if abs(cross_prod) > 1e-6: - reduced_vectors.append(p) - break - reduced_vectors = np.vstack(reduced_vectors) - return Basis2d(reduced_vectors, qmax, q_tol, theta_tol_deg) - - elif vectors.shape == (3,3): - temp_basis = Basis3d(vectors, qmax, q_tol, theta_tol_deg, symmetry, centering) - - if not reduce: # Short circuit - return temp_basis - temp_basis.generate_points_and_pairs_fast() - points = temp_basis.points - reduced_vectors = [] - distances = np.linalg.norm(points, axis=1) - sort_idx = np.argsort(distances) - - for i in sort_idx: - p = points[i] - if len(reduced_vectors) == 0 and np.linalg.norm(p) > 1e-6: - reduced_vectors.append(p) - elif len(reduced_vectors) == 1: - cross_prod = np.cross(reduced_vectors[0], p) - if np.linalg.norm(cross_prod) > 1e-6: - reduced_vectors.append(p) - elif len(reduced_vectors) == 2: - v1, v2 = reduced_vectors - det = np.dot(np.cross(v1, v2), p) - if abs(det) > 1e-6: - reduced_vectors.append(p) - break - reduced_vectors = np.vstack(reduced_vectors) - return Basis3d(reduced_vectors, qmax, q_tol, theta_tol_deg, symmetry, centering) - - @classmethod - def from_crystal_symmetry(cls, cs, **init_kwargs): - # Make cell vectors in conventional orientation - q1, q2, q3, alpha_star, beta_star, gamma_star = cs.unit_cell().reciprocal_parameters() - alpha_star_rad = np.radians(alpha_star) - beta_star_rad = np.radians(beta_star) - gamma_star_rad = np.radians(gamma_star) - # Create basis vectors using crystallographic conventions - # First vector along x - v1 = np.array([q1, 0.0, 0.0]) - - # Second vector in xy plane - v2 = q2 * np.array([np.cos(gamma_star_rad), - np.sin(gamma_star_rad), - 0.0]) - - # Third vector using all angles - cx = np.cos(beta_star_rad) - cy = (np.cos(alpha_star_rad) - - np.cos(beta_star_rad)*np.cos(gamma_star_rad))/np.sin(gamma_star_rad) - cz = np.sqrt(1.0 - cx*cx - cy*cy) - v3 = q3 * np.array([cx, cy, cz]) - - vectors = np.vstack([v1, v2, v3]) - cr_system = cs.space_group().crystal_system() - centering = cs.space_group_info().symbol_and_number()[0] - result = cls.from_vectors(vectors, reduce=False, symmetry=cr_system, - centering=centering, **init_kwargs) - return result - - - - @classmethod - def from_params(cls, *params, reduce=True, **init_kwargs): - """ - params: q1 (A-1), q2, ga* (degrees) or q1, q2, q3, al*, be*, ga*. - reduce: if True, return the setting from the shortest suitable (non-collinear - etc) vectors. If False, return the setting as given. - """ - assert len(params) in (3,6) - - if len(params) == 3: - # 2D case: q1, q2, gamma* - q1, q2, gamma_star_rad = params - - # Create basis vectors - v1 = np.array([q1, 0.0]) - v2 = q2 * np.array([np.cos(gamma_star_rad), np.sin(gamma_star_rad)]) - vectors = np.vstack([v1, v2]) - - else: - # 3D case: q1, q2, q3, alpha*, beta*, gamma* - q1, q2, q3, alpha_star, beta_star, gamma_star = params - # Convert angles to radians - alpha_star_rad = np.radians(alpha_star) - beta_star_rad = np.radians(beta_star) - gamma_star_rad = np.radians(gamma_star) - - # Create basis vectors using crystallographic conventions - # First vector along x - v1 = np.array([q1, 0.0, 0.0]) - - # Second vector in xy plane - v2 = q2 * np.array([np.cos(gamma_star_rad), - np.sin(gamma_star_rad), - 0.0]) - - # Third vector using all angles - cx = np.cos(beta_star_rad) - cy = (np.cos(alpha_star_rad) - - np.cos(beta_star_rad)*np.cos(gamma_star_rad))/np.sin(gamma_star_rad) - cz = np.sqrt(1.0 - cx*cx - cy*cy) - v3 = q3 * np.array([cx, cy, cz]) - - vectors = np.vstack([v1, v2, v3]) - - if reduce: - # Create temporary basis to generate points - temp_basis = cls.from_vectors(vectors, **init_kwargs) - try: - temp_basis.generate_points_and_pairs_fast() - except Exception: - return None - points = temp_basis.points - - if len(params) == 3: - # Find two shortest non-collinear vectors - basis_vectors = [] - distances = np.linalg.norm(points, axis=1) - sort_idx = np.argsort(distances) - - for i in sort_idx: - p = points[i] - if len(basis_vectors) == 0 and np.linalg.norm(p) > 1e-6: - basis_vectors.append(p) - elif len(basis_vectors) == 1: - cross_prod = np.cross(basis_vectors[0], p) - if abs(cross_prod) > 1e-6: - basis_vectors.append(p) - break - vectors = np.vstack(basis_vectors) - - else: - # 3D reduction (might want to implement a more sophisticated method) - # For now, just find three shortest non-coplanar vectors - basis_vectors = [] - distances = np.linalg.norm(points, axis=1) - sort_idx = np.argsort(distances) - - for i in sort_idx: - p = points[i] - if len(basis_vectors) == 0 and np.linalg.norm(p) > 1e-6: - basis_vectors.append(p) - elif len(basis_vectors) == 1: - cross_prod = np.cross(basis_vectors[0], p) - if np.linalg.norm(cross_prod) > 1e-6: - basis_vectors.append(p) - elif len(basis_vectors) == 2: - v1, v2 = basis_vectors - det = np.dot(np.cross(v1, v2), p) - if abs(det) > 1e-6: - basis_vectors.append(p) - break - vectors = np.vstack(basis_vectors) - - return cls.from_vectors(vectors, **init_kwargs) - - def match(self, pair: SpotPair) -> Tuple[Optional[dict], str]: - raise NotImplementedError - - def fom_1d(self, pairs, delta=.001): - """Calculate figure of merit as percentage of observed q-values within threshold of lattice points. - - Args: - pairs: List of SpotPair objects with observed q-values - delta: Tolerance threshold in Å-1 - - Returns: - Percentage (0-100) of observed q-values that match lattice points - """ - if not hasattr(self, 'points') or self.points is None: - self.generate_points_and_pairs_fast() - - # Extract all observed q-values - q_obs = [] - for p in pairs: - q_obs.extend([p.q1, p.q2]) - q_obs = np.array(q_obs) - - # Calculate all q-values from lattice points - q_calc = np.linalg.norm(self.points, axis=1) - - # Reshape for broadcasting - q_obs_col = q_obs.reshape(-1, 1) # shape: (n_obs, 1) - q_calc_row = q_calc.reshape(1, -1) # shape: (1, n_calc) - - # Compute differences between all pairs - diffs = np.abs(q_obs_col - q_calc_row) # shape: (n_obs, n_calc) - - has_match = np.any(diffs < delta, axis=1) - - # Compute percentage - pct_matched = 100.0 * np.sum(has_match) / len(q_obs) - - return pct_matched - - - def generate_points_and_pairs_fast(self): - """Generate unique lattice point pairs preserving q1 ordering.""" - # Generate points - max_indices = np.ceil(self.qmax / np.linalg.norm(self.vectors, axis=1)) - ranges = [np.arange(-n, n+1) for n in max_indices.astype(int)] - - # Generate mesh grid based on dimension - if len(ranges) == 2: # 2D - h_range = np.concatenate([[0], np.arange(1, max_indices[0] + 1)]) - k_range = np.arange(-max_indices[1], max_indices[1] + 1) - H, K = np.meshgrid(h_range, k_range, indexing='ij') - indices = np.column_stack((H.flatten(), K.flatten())) - else: # 3D - h_range = np.concatenate([[0], np.arange(1, max_indices[0] + 1)]) - k_range = np.arange(-max_indices[1], max_indices[1] + 1) - l_range = np.arange(-max_indices[2], max_indices[2] + 1) - H, K, L = np.meshgrid(h_range, k_range, l_range, indexing='ij') - indices = np.column_stack((H.flatten(), K.flatten(), L.flatten())) - - # Filter indices by centering type - - lattice_condition_dict = { - None: lambda x: True, - 'P': lambda x: True, - 'A': lambda x: (x[1] + x[2]) % 2 == 0, - 'B': lambda x: (x[0] + x[2]) % 2 == 0, - 'C': lambda x: (x[0] + x[1]) % 2 == 0, - 'I': lambda x: (x[0] + x[1] + x[2]) % 2 == 0, - 'F': lambda x: x[0]%2 == x[1]%2 == x[2]%2, - 'hR': lambda x: (-x[0] + x[1] + x[2]) % 3 == 0, - 'R': lambda x: (-x[0] + x[1] + x[2]) % 3 == 0, - } - lattice_condition = lattice_condition_dict[self.centering] - lattice_mask = np.array([lattice_condition(row) for row in indices]) - indices = indices[lattice_mask] - - - # Generate points - points = indices @ self.vectors - - # Filter by magnitude - magnitudes = np.linalg.norm(points, axis=1) - mask = (magnitudes <= self.qmax) & (magnitudes > 0) # Exclude origin - - # Keep only points within qmax - filtered_points = points[mask] - filtered_indices = indices[mask] - filtered_magnitudes = magnitudes[mask] - self.points = filtered_points - self.point_indices = filtered_indices - - # Sort by magnitude (this will make q1_values naturally sorted) - sort_idx = np.argsort(filtered_magnitudes) - sorted_points = filtered_points[sort_idx] - sorted_indices = filtered_indices[sort_idx] - sorted_magnitudes = filtered_magnitudes[sort_idx] - - # Create inverted versions - inverted_indices = -sorted_indices - inverted_points = inverted_indices @ self.vectors - - # Create normalized vectors for dot products - n_points = len(sorted_points) - normalized_points = sorted_points / sorted_magnitudes[:, np.newaxis] - inverted_normalized = inverted_points / np.linalg.norm(inverted_points, axis=1)[:, np.newaxis] - - # Create full matrices with NaN padding outside upper triangle - - # 1. First create mask for upper triangle (excluding diagonal) - triu_mask = np.triu(np.ones((n_points, n_points)), k=1) - nan_mask = ~triu_mask.astype(bool) - - # 2. Create q matrices (sorted by construction) - q1_matrix = np.broadcast_to(sorted_magnitudes[:, np.newaxis], (n_points, n_points)) - q2_matrix = np.broadcast_to(sorted_magnitudes[np.newaxis, :], (n_points, n_points)) - - # 3. Calculate theta matrices - dot_products_normal = np.dot(normalized_points, normalized_points.T) - theta_matrix_normal = np.arccos(np.clip(dot_products_normal, -1, 1)) - - dot_products_inverted = np.dot(normalized_points, inverted_normalized.T) - theta_matrix_inverted = np.arccos(np.clip(dot_products_inverted, -1, 1)) - - # 4. Create hkl matrices - # For normal pairs - hkl1_shape = (n_points, n_points, sorted_indices.shape[1]) - hkl1_matrix_normal = np.broadcast_to(sorted_indices[:, np.newaxis, :], hkl1_shape) - hkl2_matrix_normal = np.broadcast_to(sorted_indices[np.newaxis, :, :], hkl1_shape) - - # For inverted pairs - hkl1_matrix_inverted = hkl1_matrix_normal.copy() - hkl2_matrix_inverted = np.broadcast_to(inverted_indices[np.newaxis, :, :], hkl1_shape) - - # 5. Apply NaN mask to matrices - q1_matrix_normal = q1_matrix.copy() - q1_matrix_normal[nan_mask] = np.nan - - q2_matrix_normal = q2_matrix.copy() - q2_matrix_normal[nan_mask] = np.nan - - theta_matrix_normal[nan_mask] = np.nan - - q1_matrix_inverted = q1_matrix.copy() - q1_matrix_inverted[nan_mask] = np.nan - - q2_matrix_inverted = q2_matrix.copy() - q2_matrix_inverted[nan_mask] = np.nan - - theta_matrix_inverted[nan_mask] = np.nan - - # 6. Stack normal and inverted matrices horizontally - q1_stacked = np.hstack([q1_matrix_normal, q1_matrix_inverted]) - q2_stacked = np.hstack([q2_matrix_normal, q2_matrix_inverted]) - theta_stacked = np.hstack([theta_matrix_normal, theta_matrix_inverted]) - - # 7. Flatten stacked matrices - q1_flat = q1_stacked.flatten() - q2_flat = q2_stacked.flatten() - theta_flat = theta_stacked.flatten() - - # 8. Remove NaN values - valid_mask = ~np.isnan(q1_flat) - self.q1_values = q1_flat[valid_mask] - self.q2_values = q2_flat[valid_mask] - self.theta_values = theta_flat[valid_mask] - - # Handle hkl values (more complex due to extra dimension) - hkl_dim = sorted_indices.shape[1] - - # Apply NaN mask to hkl matrices via boolean indexing - # Create "is NaN" array matching hkl shape - nan_mask_expanded = np.broadcast_to(nan_mask[:, :, np.newaxis], - (n_points, n_points, hkl_dim)) - - # Set invalid hkls to a dummy value (will be filtered out later) - hkl1_matrix_normal = np.where(nan_mask_expanded, -999, hkl1_matrix_normal) - hkl2_matrix_normal = np.where(nan_mask_expanded, -999, hkl2_matrix_normal) - hkl1_matrix_inverted = np.where(nan_mask_expanded, -999, hkl1_matrix_inverted) - hkl2_matrix_inverted = np.where(nan_mask_expanded, -999, hkl2_matrix_inverted) - - # Stack hkl matrices horizontally - hkl1_stacked = np.hstack([hkl1_matrix_normal.reshape(n_points, -1), - hkl1_matrix_inverted.reshape(n_points, -1)]) - hkl2_stacked = np.hstack([hkl2_matrix_normal.reshape(n_points, -1), - hkl2_matrix_inverted.reshape(n_points, -1)]) - - # Reshape to get original structure back - hkl1_flat = hkl1_stacked.reshape(-1, hkl_dim) - hkl2_flat = hkl2_stacked.reshape(-1, hkl_dim) - - # Apply same valid mask as for q values - self.hkl1_values = hkl1_flat[valid_mask] - self.hkl2_values = hkl2_flat[valid_mask] - - # Ensure q1 <= q2 (swap if needed) - swap_mask = self.q1_values > self.q2_values - if np.any(swap_mask): - print('swap') - self.q1_values[swap_mask], self.q2_values[swap_mask] = self.q2_values[swap_mask], self.q1_values[swap_mask].copy() - self.hkl1_values[swap_mask], self.hkl2_values[swap_mask] = self.hkl2_values[swap_mask], self.hkl1_values[swap_mask].copy() - def match_pairs(self, pairs): - self.all_pairs = [] - for p in pairs: - result, status = self.match(p) - self.all_pairs.append((p, result, status)) - - def reindex_pairs(self): - self.match_pairs([p[0] for p in self.all_pairs]) - - - - - - -class Basis2d(Basis): - - def index_percent(self): - all = 0 - hits = 0 - for p in self.all_pairs: - all += 1 - if p[2]=='indexed_2d': - hits += 1 - return hits/all - def fom_2d(self, pairs): - hits = 0 - for p in pairs: - result = self.match(p) - if result[1]=='indexed_2d': - hits += 1 - return 100 * hits/len(pairs) - def match(self, pair: SpotPair) -> Tuple[Optional[dict], str]: - - if not hasattr(self, 'q1_values'): - self.generate_points_and_pairs_fast() - - # Compute differences with input pair - dq1 = np.abs(self.q1_values - pair.q1) - dq2 = np.abs(self.q2_values - pair.q2) - dtheta = np.abs(self.theta_values - pair.theta) - - # Create mask for matches within tolerance - matches = (dq1 < self.q_tolerance) & (dq2 < self.q_tolerance) & (dtheta < self.theta_tolerance) - - if np.any(matches): - # Get the first match (or could find best one later) - idx = np.where(matches)[0][0] - hkl1 = self.hkl1_values[idx] - hkl2 = self.hkl2_values[idx] - - result = PairMatch2d(hkl1, pair.q1, hkl2, pair.q2, pair.theta) - return result, 'indexed_2d' - - qmags = np.linalg.norm(self.points, axis=1) - dq1 = np.abs(pair.q1 - qmags) - dq2 = np.abs(pair.q2 - qmags) - i_best = np.argmin(np.minimum(dq1, dq2)) - one_vec_match = False - if dq1[i_best] < self.q_tolerance: - one_vec_match = True - q1 = pair.q1 - q2 = pair.q2 - elif dq2[i_best] < self.q_tolerance: - one_vec_match = True - q1 = pair.q2 - q2 = pair.q1 - if one_vec_match: - return OneVectorMatch(self.point_indices[i_best], q1, q2, pair.theta), 'one_vector' - return None, 'unindexed' - - def flag_outliers(self, multiplier=2, q_weight=1000, theta_weight=111): - self.pairs_costs = [] - for p in self.all_pairs: - if p[2]=='indexed_2d': - self.pairs_costs.append(( - p, - self.compute_pair_cost(self.vectors, p[1], q_weight, theta_weight)) - ) - max_delta = np.median([pc[1] for pc in self.pairs_costs]) * multiplier - for p, c in self.pairs_costs: - p[1].is_outlier = c > max_delta - def plot_costs(self, q_weight=1000, theta_weight=111, cost_max=2): - costs = [pc[1] for pc in self.pairs_costs] - costs_inlier = [pc[1] for pc in self.pairs_costs if not pc[0][1].is_outlier] - costs_outlier = [pc[1] for pc in self.pairs_costs if pc[0][1].is_outlier] -# for p in self.sublattice_indexed: -# costs.append(self.sub_basis.compute_pair_cost( -# self.sub_basis.vectors, p, q_weight, theta_weight, use_outliers=True)) -# costs_inlier = [c for c,p in zip(costs, self.sublattice_indexed) if not p.is_outlier] -# costs_outlier = [c for c,p in zip(costs, self.sublattice_indexed) if p.is_outlier] - bins = np.linspace(0, cost_max, 50) - plt.hist(costs_inlier, bins=bins) - plt.hist(costs_outlier, bins=bins, color='red') - plt.show() - - - - def astar(self): - return(np.linalg.norm(self.vectors[0])) - def bstar(self): - return(np.linalg.norm(self.vectors[1])) - def gammastar(self): - result = np.degrees(np.arccos( - np.dot(self.vectors[0], self.vectors[1]) / (a * b))) - return result - - def __str__(self): - """Format 2D sublattice parameters.""" - # Calculate a, b, gamma from basis vectors - a = np.linalg.norm(self.vectors[0]) - b = np.linalg.norm(self.vectors[1]) - gamma = np.degrees(np.arccos( - np.dot(self.vectors[0], self.vectors[1]) / (a * b))) - - return f"a={a:.5f}, b={b:.5f}, gamma={gamma:.2f}°" - - def area(self): - return abs(np.cross(*self.vectors)) - - def doubled_cells(self) -> list: - """Generate all doubled and tripled variants of the current 2D basis. - - Returns: - List of 7 Basis2d objects: 3 doubled cells followed by 4 tripled cells - """ - # Get current reciprocal basis vectors - a_star = self.vectors[0].copy() - b_star = self.vectors[1].copy() - - result = [] - - # DOUBLED CELLS (3 cases) - # 1. Double a: a*' = a*/2, b*' = b* - basis1 = self.vectors.copy() - basis1[0] = a_star / 2 - result.append(basis1) - - # 2. Double b: a*' = a*, b*' = b*/2 - basis2 = self.vectors.copy() - basis2[1] = b_star / 2 - result.append(basis2) - - # 3. Double along (1,1): a*' = (a*+b*)/2, b*' = (a*-b*)/2 - basis3 = self.vectors.copy() - basis3[0] = (a_star + b_star) / 2 - basis3[1] = (a_star - b_star) / 2 - result.append(basis3) - - # TRIPLED CELLS (4 cases) - # 4. Triple a: a*' = a*/3, b*' = b* - basis4 = self.vectors.copy() - basis4[0] = a_star / 3 - result.append(basis4) - - # 5. Triple b: a*' = a*, b*' = b*/3 - basis5 = self.vectors.copy() - basis5[1] = b_star / 3 - result.append(basis5) - - # 6. Triple along (1,1): points at (1/3,1/3) and (2/3,-1/3) - basis6 = self.vectors.copy() - basis6[0] = (a_star + b_star) / 3 - basis6[1] = (2*a_star - b_star) / 3 - result.append(basis6) - - # 7. Triple along (1,-1): points at (1/3,-1/3) and (2/3,1/3) - basis7 = self.vectors.copy() - basis7[0] = (a_star - b_star) / 3 - basis7[1] = (2*a_star + b_star) / 3 - result.append(basis7) - - # Create new Basis2d objects - final_result = [] - for basis in result: - new_basis = type(self).from_vectors( - basis, - qmax=self.qmax, - q_tol=self.q_tolerance, - theta_tol_deg=np.degrees(self.theta_tolerance) - ) - final_result.append(new_basis) - - return final_result - - def compute_pair_cost(self, basis_vectors: np.ndarray, pair: PairMatch2d, - q_weight: float = 1.0, theta_weight: float = 111.0, use_outliers=False) -> float: - """Compute cost for a single indexed pair using Miller indices.""" - if pair.is_outlier and not use_outliers: return 0 - # Compute q-vectors from Miller indices - q1_calc = pair.hkl1[0]*basis_vectors[0] + pair.hkl1[1]*basis_vectors[1] - q2_calc = pair.hkl2[0]*basis_vectors[0] + pair.hkl2[1]*basis_vectors[1] - - # Compare magnitudes - q1_calc_mag = np.linalg.norm(q1_calc) - q2_calc_mag = np.linalg.norm(q2_calc) - q_error = (abs(q1_calc_mag - pair.q1) + - abs(q2_calc_mag - pair.q2)) - - # Compare angle - cos_theta_calc = np.dot(q1_calc, q2_calc) / (q1_calc_mag * q2_calc_mag) - theta_calc = np.arccos(np.clip(cos_theta_calc, -1, 1)) - theta_error = abs(theta_calc - pair.theta_rad) - - return q_weight * q_error + theta_weight * theta_error - - def cost(self, params, pairs, q_weight=1000, theta_weight=111.0): - """Total cost for given lattice parameters.""" - a_star, b_star, gamma_star = params - - # Basis vectors in standard orientation - v1 = np.array([a_star, 0.0]) - v2 = b_star * np.array([np.cos(gamma_star), np.sin(gamma_star)]) - basis_vectors = np.vstack([v1, v2]) - - # Sum costs from all pairs - return sum(self.compute_pair_cost(basis_vectors, pair, q_weight, theta_weight) - for pair in pairs) - - def plot_cost_histogram(self, pairs=None, q_weight=1000, theta_weight=111.0): - - if pairs is None: - pairs = [p[1] for p in self.all_pairs if p[2]=='indexed_2d'] - #copied from refine, sue me - a_star = np.linalg.norm(self.vectors[0]) - b_star = np.linalg.norm(self.vectors[1]) - gamma_star = np.arccos(np.dot(self.vectors[0], self.vectors[1]) / - (a_star * b_star)) - v1 = np.array([a_star, 0.0]) - v2 = b_star * np.array([np.cos(gamma_star), np.sin(gamma_star)]) - basis_vectors = np.vstack([v1, v2]) - - costs = [self.compute_pair_cost(basis_vectors, p, q_weight, theta_weight) for p in pairs] - plt.hist(costs, bins=100) - plt.show() - - def refine(self, pairs=None, q_weight: float = 1000.0, theta_weight: float = 111.0) -> None: - """Refine lattice parameters in place.""" - - if pairs is None: - pairs = [p[1] for p in self.all_pairs if p[2]=='indexed_2d'] - - # Initial parameter vector - a_star = np.linalg.norm(self.vectors[0]) - b_star = np.linalg.norm(self.vectors[1]) - gamma_star = np.arccos(np.dot(self.vectors[0], self.vectors[1]) / - (a_star * b_star)) - initial_params = np.array([a_star, b_star, gamma_star]) - ip = copy.deepcopy(initial_params) - - # Run optimization - result = minimize( - lambda p: self.cost(p, pairs, q_weight, theta_weight), - initial_params, - method='BFGS', - #options={'gtol': 1e-8} - ) - - if not result.success: - print(result.message) - - # Extract refined parameters and update basis - a_star, b_star, gamma_star = result.x - if gamma_star < np.pi/2: gamma_star = np.pi - gamma_star - self.vectors = np.array([[a_star, 0.0], - [b_star * np.cos(gamma_star), - b_star * np.sin(gamma_star)]]) - - # Regenerate points and pairs with new basis - self.generate_points_and_pairs_fast() - - - - -class Basis3d(Basis): - - def index_percent(self): - all = 0 - hits = 0 - for p in self.all_pairs: - all += 1 - if p[2]=='indexed_3d': - hits += 1 - return hits/all - - def flag_outliers(self, multiplier=2, q_weight=1000, theta_weight=111): - self.pairs_costs = [] - for p in self.all_pairs: - if p[2]=='indexed_3d': - self.pairs_costs.append(( - p, - self.compute_pair_cost(self.vectors, p[1], q_weight, theta_weight)) - ) - max_delta = np.median([pc[1] for pc in self.pairs_costs]) * multiplier - for p, c in self.pairs_costs: - p[1].is_outlier = c > max_delta - - def plot_costs(self, q_weight=1000, theta_weight=111): - costs = [pc[1] for pc in self.pairs_costs] - costs_inlier = [pc[1] for pc in self.pairs_costs if not pc[0][1].is_outlier] - costs_outlier = [pc[1] for pc in self.pairs_costs if pc[0][1].is_outlier] -# for p in self.sublattice_indexed: -# costs.append(self.sub_basis.compute_pair_cost( -# self.sub_basis.vectors, p, q_weight, theta_weight, use_outliers=True)) -# costs_inlier = [c for c,p in zip(costs, self.sublattice_indexed) if not p.is_outlier] -# costs_outlier = [c for c,p in zip(costs, self.sublattice_indexed) if p.is_outlier] - bins = np.linspace(0, 2, 50) - plt.hist(costs_inlier, bins=bins) - plt.hist(costs_outlier, bins=bins, color='red') - plt.show() - - def plot_costs_components(self, q_weight=1000, theta_weight=111): - theta_costs, q_costs = [],[] - for p in self.all_pairs: - if p[2]=='indexed_3d': - theta_costs.append( - self.compute_pair_cost(self.vectors, p[1], q_weight=0, theta_weight=theta_weight) - ) - q_costs.append( - self.compute_pair_cost(self.vectors, p[1], q_weight, theta_weight=0) - ) - bins = np.linspace(0, 2, 50) - fig, (ax1, ax2) = plt.subplots(2,1) - ax1.set_title('Costs (theta component)') - ax1.hist(theta_costs, bins=bins) - ax1.xaxis.set_visible(False) - ax2.set_title('Costs (q component)') - ax2.hist(q_costs, bins=bins) - plt.show() - - - def volume(self) -> float: - """Calculate the volume of the reciprocal space unit cell. - - Returns: - Volume in Å⁻³ - """ - return np.abs(np.linalg.det(self.vectors)) - - def __str__(self): - # Reciprocal cell parameters - a_star = np.linalg.norm(self.vectors[0]) - b_star = np.linalg.norm(self.vectors[1]) - c_star = np.linalg.norm(self.vectors[2]) - - alpha_star = np.degrees(np.arccos( - np.dot(self.vectors[1], self.vectors[2]) / (b_star * c_star))) - beta_star = np.degrees(np.arccos( - np.dot(self.vectors[0], self.vectors[2]) / (a_star * c_star))) - gamma_star = np.degrees(np.arccos( - np.dot(self.vectors[0], self.vectors[1]) / (a_star * b_star))) - - # Direct cell parameters - direct = self.compute_direct_cell_params(self.vectors) - - return \ - str(round(1/self.volume(), 2)) + ' ' + ','.join( [str(round(direct[x], 3)) for x in ['a','b','c','alpha','beta','gamma']] ) + ' ' + self.centering - - - def compute_direct_cell_params(self, recip_basis: np.ndarray) -> dict: - """Compute direct cell parameters from reciprocal space basis vectors. - - Args: - recip_basis: 3x3 matrix of reciprocal space basis vectors - - Returns: - dict with a,b,c (in Å) and alpha,beta,gamma (in degrees) - """ - # Compute reciprocal metric tensor G* = B B^T - G_star = recip_basis @ recip_basis.T - - # Invert to get direct metric tensor G = (G*)^-1 - G = np.linalg.inv(G_star) - - # Extract cell parameters - a = np.sqrt(G[0,0]) - b = np.sqrt(G[1,1]) - c = np.sqrt(G[2,2]) - - alpha = np.degrees(np.arccos(G[1,2] / (b*c))) - beta = np.degrees(np.arccos(G[0,2] / (a*c))) - gamma = np.degrees(np.arccos(G[0,1] / (a*b))) - - return { - 'a': a, - 'b': b, - 'c': c, - 'alpha': alpha, - 'beta': beta, - 'gamma': gamma - } - - def match(self, pair: SpotPair) -> Tuple[Optional[dict], str]: - """Find a matching pair in the lattice using binary search on sorted values.""" - - if not hasattr(self, 'q1_values'): - self.generate_points_and_pairs_fast() - - # Use binary search to find range of indices where q1 is within tolerance - q1_min = pair.q1 - self.q_tolerance - q1_max = pair.q1 + self.q_tolerance - - # Find indices where q1_values are in range [q1_min, q1_max] - left_idx = np.searchsorted(self.q1_values, q1_min, side='left') - right_idx = np.searchsorted(self.q1_values, q1_max, side='right') - - # If no values in range, return unindexed - if left_idx >= right_idx: - return None, 'unindexed' - - # Get the subset of potential matches - subset_slice = slice(left_idx, right_idx) - q2_subset = self.q2_values[subset_slice] - theta_subset = self.theta_values[subset_slice] - - # Check q2 and theta matches in the smaller subset - match_dq2 = (q2_subset > pair.q2 - self.q_tolerance) & \ - (q2_subset < pair.q2 + self.q_tolerance) - match_theta = (theta_subset > pair.theta - self.theta_tolerance) & \ - (theta_subset < pair.theta + self.theta_tolerance) - - matches = match_dq2 & match_theta - - if np.any(matches): - # Get the first match - match_idx = np.where(matches)[0][0] - full_idx = left_idx + match_idx # Adjust back to original array index - - hkl1 = self.hkl1_values[full_idx] - hkl2 = self.hkl2_values[full_idx] - - result = PairMatch2d(hkl1, pair.q1, hkl2, pair.q2, pair.theta) - return result, 'indexed_3d' - - return None, 'unindexed' - - def doubled_cells(self) -> list: - """Generate 6 cell-doubled variants of the current basis. - - Returns: - List of 6 Basis3d objects with doubled cells: - [0] Double a: a'=2a, b'=b, c'=c - [1] Double b: a'=a, b'=2b, c'=c - [2] Double c: a'=a, b'=b, c'=2c - [3] Double ab plane: a'=a+b, b'=a-b, c'=c - [4] Double ac plane: a'=a+c, b'=b, c'=a-c - [5] Double bc plane: a'=a, b'=b+c, c'=b-c - """ - result = [] - - # Get current reciprocal basis vectors - a_star = self.vectors[0].copy() - b_star = self.vectors[1].copy() - c_star = self.vectors[2].copy() - - # Create transformation matrices in reciprocal space - # When we double a direct space vector, the corresponding reciprocal vector is halved - - # 1. Double a: a*' = a*/2, b*' = b*, c*' = c* - basis1 = self.vectors.copy() - basis1[0] = a_star / 2 - result.append(basis1) - - # 2. Double b: a*' = a*, b*' = b*/2, c*' = c* - basis2 = self.vectors.copy() - basis2[1] = b_star / 2 - result.append(basis2) - - # 3. Double c: a*' = a*, b*' = b*, c*' = c*/2 - basis3 = self.vectors.copy() - basis3[2] = c_star / 2 - result.append(basis3) - - # 4. Double ab plane: a'=a+b, b'=a-b in direct space - # In reciprocal space: a*' = (a*+b*)/2, b*' = (a*-b*)/2, c*' = c* - basis4 = self.vectors.copy() - basis4[0] = (a_star + b_star) / 2 - basis4[1] = (a_star - b_star) / 2 - result.append(basis4) - - # 5. Double ac plane: a'=a+c, b'=b, c'=a-c in direct space - # In reciprocal space: a*' = (a*+c*)/2, b*' = b*, c*' = (a*-c*)/2 - basis5 = self.vectors.copy() - basis5[0] = (a_star + c_star) / 2 - basis5[2] = (a_star - c_star) / 2 - result.append(basis5) - - # 6. Double bc plane: a'=a, b'=b+c, c'=b-c in direct space - # In reciprocal space: a*' = a*, b*' = (b*+c*)/2, c*' = (b*-c*)/2 - basis6 = self.vectors.copy() - basis6[1] = (b_star + c_star) / 2 - basis6[2] = (b_star - c_star) / 2 - result.append(basis6) - - # 7. Body-centered (I-centered) - # In reciprocal space: a*'=a*, b*'=b*, c*'=a*/2+b*/2+c*/2 - basis7 = self.vectors.copy() - basis7[2] = (a_star + b_star + c_star) / 2 - result.append(basis7) - - # Now the tripled cells - - # Type 1: Along coordinate axes (3 cases) - # 1. Triple along a: a*' = a*/3 - basis1 = self.vectors.copy() - basis1[0] = a_star / 3 - result.append(basis1) - - # 2. Triple along b: b*' = b*/3 - basis2 = self.vectors.copy() - basis2[1] = b_star / 3 - result.append(basis2) - - # 3. Triple along c: c*' = c*/3 - basis3 = self.vectors.copy() - basis3[2] = c_star / 3 - result.append(basis3) - - # Type 2: Along face diagonals (6 cases) - # 4. Along (1,1,0): points (1/3,1/3,0) and (2/3,-1/3,0) - basis4 = self.vectors.copy() - basis4[0] = (a_star + b_star) / 3 - basis4[1] = (2*a_star - b_star) / 3 - result.append(basis4) - - # 5. Along (1,-1,0): points (1/3,-1/3,0) and (2/3,1/3,0) - basis5 = self.vectors.copy() - basis5[0] = (a_star - b_star) / 3 - basis5[1] = (2*a_star + b_star) / 3 - result.append(basis5) - - # 6. Along (1,0,1): points (1/3,0,1/3) and (2/3,0,-1/3) - basis6 = self.vectors.copy() - basis6[0] = (a_star + c_star) / 3 - basis6[2] = (2*a_star - c_star) / 3 - result.append(basis6) - - # 7. Along (1,0,-1): points (1/3,0,-1/3) and (2/3,0,1/3) - basis7 = self.vectors.copy() - basis7[0] = (a_star - c_star) / 3 - basis7[2] = (2*a_star + c_star) / 3 - result.append(basis7) - - # 8. Along (0,1,1): points (0,1/3,1/3) and (0,2/3,-1/3) - basis8 = self.vectors.copy() - basis8[1] = (b_star + c_star) / 3 - basis8[2] = (2*b_star - c_star) / 3 - result.append(basis8) - - # 9. Along (0,1,-1): points (0,1/3,-1/3) and (0,2/3,1/3) - basis9 = self.vectors.copy() - basis9[1] = (b_star - c_star) / 3 - basis9[2] = (2*b_star + c_star) / 3 - result.append(basis9) - - # Type 3: Along body diagonals (4 cases) - # Group 8 (body diagonal case 1) - basis10 = self.vectors.copy() - basis10[0] = (a_star - 2*b_star - 2*c_star) / 3 - basis10[1] = b_star - basis10[2] = c_star - result.append(basis10) - - # Group 9 (body diagonal case 2) - basis11 = self.vectors.copy() - basis11[0] = (a_star - b_star - 2*c_star) / 3 - basis11[1] = b_star - basis11[2] = c_star - result.append(basis11) - - # Group 10 (body diagonal case 3) - basis12 = self.vectors.copy() - basis12[0] = (a_star - 2*b_star - c_star) / 3 - basis12[1] = b_star - basis12[2] = c_star - result.append(basis12) - - # Group 11 (body diagonal case 4) - basis13 = self.vectors.copy() - basis13[0] = (a_star - b_star - c_star) / 3 - basis13[1] = b_star - basis13[2] = c_star - result.append(basis13) - - - # Create new Basis3d objects - final_result = [] - for basis in result: - new_basis = type(self).from_vectors( - basis, - qmax=self.qmax, - q_tol=self.q_tolerance, - theta_tol_deg=np.degrees(self.theta_tolerance) - ) - final_result.append(new_basis) - - - return final_result - - def compute_pair_cost(self, basis_vectors: np.ndarray, pair: PairMatch2d, - q_weight: float = 1.0, theta_weight: float = 111.0, use_outliers=False) -> float: - """Compute cost for a single indexed pair using Miller indices.""" - if not use_outliers and pair.is_outlier: return 0 - # Compute q-vectors from Miller indices - q1_calc = pair.hkl1[0]*basis_vectors[0] + pair.hkl1[1]*basis_vectors[1] + pair.hkl1[2]*basis_vectors[2] - q2_calc = pair.hkl2[0]*basis_vectors[0] + pair.hkl2[1]*basis_vectors[1] + pair.hkl2[2]*basis_vectors[2] - - # Compare magnitudes - q1_calc_mag = np.linalg.norm(q1_calc) - q2_calc_mag = np.linalg.norm(q2_calc) - q_error = (abs(q1_calc_mag - pair.q1) + - abs(q2_calc_mag - pair.q2)) - - # Compare angle - cos_theta_calc = np.dot(q1_calc, q2_calc) / (q1_calc_mag * q2_calc_mag) - theta_calc = np.arccos(np.clip(cos_theta_calc, -1, 1)) - theta_error = abs(theta_calc - pair.theta_rad) - - return q_weight * q_error + theta_weight * theta_error - - def vectors_from_params(self, params): - """Create basis vectors from refinement parameters.""" - a_star, b_star, c_star, alpha_star, beta_star, gamma_star = params - - # First vector along x - v1 = np.array([a_star, 0.0, 0.0]) - - # Second vector in xy plane - v2 = b_star * np.array([np.cos(gamma_star), - np.sin(gamma_star), - 0.0]) - - # Third vector using all angles - cx = np.cos(beta_star) - cy = (np.cos(alpha_star) - - np.cos(beta_star)*np.cos(gamma_star))/np.sin(gamma_star) - cz_sq = 1.0 - cx*cx - cy*cy - if cz_sq < 0: - # This happens with invalid angle combinations - cz = 0.0 - # Add large penalty to cost function - return None - else: - cz = np.sqrt(cz_sq) - v3 = c_star * np.array([cx, cy, cz]) - - return np.vstack([v1, v2, v3]) - - def cost(self, params, pairs, q_weight=1000, theta_weight=111.0): - """Total cost for given lattice parameters.""" - # Generate basis vectors from parameters - basis_vectors = self.vectors_from_params(params) - - # If invalid parameters, return large penalty - if basis_vectors is None: - return 1.0e6 - - # Sum costs from all pairs - return sum(self.compute_pair_cost(basis_vectors, pair, q_weight, theta_weight) - for pair in pairs) - - def plot_cost_histogram(self, pairs, q_weight=1000, theta_weight=111.0): - - #copied from refine, sue me - a_star = np.linalg.norm(self.vectors[0]) - b_star = np.linalg.norm(self.vectors[1]) - c_star = np.linalg.norm(self.vectors[2]) - - alpha_star = np.arccos(np.dot(self.vectors[1], self.vectors[2]) / - (b_star * c_star)) - beta_star = np.arccos(np.dot(self.vectors[0], self.vectors[2]) / - (a_star * c_star)) - gamma_star = np.arccos(np.dot(self.vectors[0], self.vectors[1]) / - (a_star * b_star)) - - current_params = np.array([a_star, b_star, c_star, - alpha_star, beta_star, gamma_star]) - current_vectors = self.vectors_from_params(current_params) - costs = [self.compute_pair_cost(current_params, p, q_weight, theta_weight) for p in pairs] - plt.hist(costs, bins=100) - plt.show() - - def refine(self, pairs=None, q_weight: float = 1000.0, theta_weight: float = 111.0, constrained=False) -> None: - """Refine lattice parameters in place.""" - - if pairs is None: - pairs = [p[1] for p in self.all_pairs if p[2]=='indexed_3d'] - - # Initial parameter vector from current basis - a_star = np.linalg.norm(self.vectors[0]) - b_star = np.linalg.norm(self.vectors[1]) - c_star = np.linalg.norm(self.vectors[2]) - - alpha_star = np.arccos(np.dot(self.vectors[1], self.vectors[2]) / - (b_star * c_star)) - beta_star = np.arccos(np.dot(self.vectors[0], self.vectors[2]) / - (a_star * c_star)) - gamma_star = np.arccos(np.dot(self.vectors[0], self.vectors[1]) / - (a_star * b_star)) - - if constrained: - assert self.symmetry in ['Triclinic', 'Monoclinic', 'Orthorhombic', - 'Tetragonal', 'Trigonal', 'Hexagonal', 'Cubic'] - - if not constrained or self.symmetry=='Triclinic': - initial_params = np.array([a_star, b_star, c_star, - alpha_star, beta_star, gamma_star]) - bounds = [ - (0.01, None), # a_star > 0 - (0.01, None), # b_star > 0 - (0.01, None), # c_star > 0 - (0.1, np.pi-0.1), # alpha_star in (0, pi) - (0.1, np.pi-0.1), # beta_star in (0, pi) - (0.1, np.pi-0.1) # gamma_star in (0, pi) - ] - def params_to_full(p): - return p - - elif self.symmetry == 'Monoclinic': - initial_params = np.array([a_star, b_star, c_star, beta_star]) - bounds = [ - (0.01, None), # a_star > 0 - (0.01, None), # b_star > 0 - (0.01, None), # c_star > 0 - (0.1, np.pi-0.1) # beta_star in (0, pi) - ] - def params_to_full(p): - a_star, b_star, c_star, beta_star = p - return np.array([a_star, b_star, c_star, np.pi/2, beta_star, np.pi/2]) - - elif self.symmetry == 'Orthorhombic': - initial_params = np.array([a_star, b_star, c_star]) - bounds = [ - (0.01, None), # a_star > 0 - (0.01, None), # b_star > 0 - (0.01, None), # c_star > 0 - ] - def params_to_full(p): - a_star, b_star, c_star = p - return np.array([a_star, b_star, c_star, np.pi/2, np.pi/2, np.pi/2]) - elif self.symmetry == 'Tetragonal': - initial_params = np.array([a_star, c_star]) - bounds = [ - (0.01, None), # a_star > 0 - (0.01, None), # b_star > 0 - ] - def params_to_full(p): - a_star, c_star = p - return np.array([a_star, a_star, c_star, np.pi/2, np.pi/2, np.pi/2]) - elif self.symmetry in ['Trigonal', 'Hexagonal']: - initial_params = np.array([a_star, c_star]) - bounds = [ - (0.01, None), # a_star > 0 - (0.01, None), # b_star > 0 - ] - def params_to_full(p): - a_star, c_star = p - return np.array([a_star, a_star, c_star, np.pi/2, np.pi/2, np.pi/3]) - elif self.symmetry == 'Cubic': - initial_params = np.array([a_star]) - bounds = [ - (0.01, None), # a_star > 0 - ] - def params_to_full(p): - a_star = p[0] - return np.array([a_star, a_star, a_star, np.pi/2, np.pi/2, np.pi/2]) - else: - raise RuntimeError('Unknown crystal system') - - def cost_wrapper(reduced_params): - full_params = params_to_full(reduced_params) - return self.cost(full_params, pairs, q_weight, theta_weight) - - # Run optimization - result = minimize( - cost_wrapper, - initial_params, - method='L-BFGS-B', # Use bounded optimization - bounds=bounds, - options={'ftol': 1e-10} - ) - - if not result.success: - print(f"Refinement warning: {result.message}") - - # Extract refined parameters and update basis - refined_params = params_to_full(result.x) - refined_vectors = self.vectors_from_params(refined_params) - - if refined_vectors is not None: - self.vectors = refined_vectors - - # Regenerate points and pairs with new basis - self.generate_points_and_pairs_fast() - - # Print refinement results - refined_vols = self.volume() - #print(f"Refinement complete. Final cost: {result.fun:.6f}") - #print(f"Reciprocal volume: {refined_vols:.6f} Å⁻³") - #print(f"Direct cell volume: {1/refined_vols:.1f} ų") - print(f"Refine done: {self}") - print(f"Volume: {1/refined_vols:.2f}") - - -class LatticeReconstruction: - def __init__(self, qmax: float, q_tolerance: float = 0.001, - theta_tol_degrees: float = 1): - self.qmax = qmax - self.q_tolerance = q_tolerance - self.theta_tol_degrees = theta_tol_degrees - - self.all_pairs = [] - self.sublattice_indexed = [] - self.sublattice_1vector = [] - self.unindexed = [] - self.current_basis = None - self.sub_basis = None - - def calculate_sublattice_area(self) -> float: - """Calculate area of current sublattice""" - if self.sub_basis is None: - return float('inf') - return np.abs(np.linalg.det(self.sub_basis)) - - def read_triplet(self, q1: float, q2: float, theta: float) -> None: - """Read a new q1,q2,theta triplet and update reconstruction.""" - pair = SpotPair(q1, q2, np.radians(theta)) # assume input theta in degrees - self.all_pairs.append(pair) - self.update() - - def store_triplet(self, q1, q2, theta): - pair = SpotPair(q1, q2, np.radians(theta)) - self.all_pairs.append(pair) - - def generate_2d_bases(self, scan_pts=11, scan_range=1, max_axis=40): - scan_range=np.radians(scan_range) - self.basis_candidates_2d = [] - for pair in self.all_pairs: - deltas = np.linspace(-scan_range, scan_range, scan_pts) - for d in deltas: - try: - b = Basis.from_params(pair.q1, pair.q2, pair.theta+d) - if 1/b.astar() < max_axis and 1/b.bstar() < max_axis: - self.basis_candidates_2d.append(b) - except Exception: - pass - - - def update(self) -> None: - """Main update method implementing the algorithm.""" - pair = self.all_pairs[-1] # Most recently read - - # Initialize first sublattice - if len(self.all_pairs) == 1: - self._initialize_first_sublattice(pair) - self.print_status() - return - - # Next, test for smaller sublattice - elif pair.area() < self.sub_basis.area(): - trial_basis = Basis.from_params(pair.q1, pair.q2, pair.theta) - print(f'Found smaller 2d basis: {trial_basis}') - print(f'From pair {pair.q1}, {pair.q2}, {np.degrees(pair.theta)}') - reproc = input('Reprocess? y/[n] ')=='y' - if reproc: - self.current_basis = None # Wipe out any 3d basis - self._initialize_first_sublattice(pair) - self.reprocess_all_pairs() - else: - pass - - # Try 3d indexing if possible - elif self.current_basis is not None: - result, status = self.current_basis.match(pair) - if status=='indexed': - self._store_result(result, status) - self.print_status() - return - - - # Finally, try matching in the current sublattice - else: - - # Process new pair - result, status = self.sub_basis.match(pair) - self._store_result(result, status) - - self.print_status() - - def _initialize_first_sublattice(self, pair: SpotPair) -> None: - """Set up initial 2D sublattice from first pair using least oblique cell.""" - self.sub_basis = Basis.from_params(pair.q1, pair.q2, pair.theta) - self.sl_from_pair = pair - - def set_sub_basis(self, basis): - self.sub_basis = basis - self.reprocess_all_pairs() - - def summarize_half_indexed_pairs(self) -> str: - """Create summary table of half-indexed pairs.""" - if not self.sublattice_1vector: - return "No half-indexed pairs found." - - lines = ["Half-indexed pairs for sublattice:", - "i | q[uni] | q[idx] | theta | err [Å⁻¹]"] - entries = [] - - for i, pair in enumerate(self.sublattice_1vector): - # For each pair, determine which q is indexed - q_idx, q_uni = pair.q1, pair.q2 - - # Find nearest sublattice point to the indexed vector - diffs = np.abs(np.linalg.norm(self.sub_basis.points, axis=1) - q_idx) - err = np.min(diffs) - - #lines.append(f"{i:2d} {q_uni:.4f} {q_idx:.4f} {np.degrees(pair.theta_rad):.1f} {err:.4f}") - entries.append((i, q_uni, q_idx, pair.theta_rad, err)) - - for i, q_uni, q_idx, theta_rad, err in sorted( - entries, - #key=lambda x:(round(x[1], 3), x[4]) - key=lambda x:(round(x[4], 4), round(x[1], 3)) - )[:50]: - lines.append(f"{i:2d}\t{q_uni:.4f}\t{q_idx:.4f}\t{np.degrees(theta_rad):.1f}\t{err:.4f}") - result = "\n".join(lines) - - return result - - def find_matching_pairs(self) -> List[Tuple[int, int, float]]: - """Find pairs of half-indexed vectors with matching unindexed q values. - - Returns: - List of (i1, i2, q) tuples where: - i1, i2: indices into sublattice_1vector - q: the matching q value - """ - matches = [] - n = len(self.sublattice_1vector) - - for i in range(n): - pair_i = self.sublattice_1vector[i] - # Get unindexed q value - q_i_uni = pair_i.q2 # The unindexed vector - q_i_idx = pair_i.q1 # The indexed vector - #q_i = pair_i.q2 if np.array_equal(pair_i.indexed_vector, pair_i.v1) else pair_i.q1 - - for j in range(i+1, n): - pair_j = self.sublattice_1vector[j] - # Get unindexed q value - q_j_uni = pair_j.q2 - q_j_idx = pair_j.q1 - - if abs(q_i_uni - q_j_uni) < self.q_tolerance and abs(q_i_idx - q_j_idx) > self.q_tolerance: - # Use average q value for the match - q_match = (q_i_uni + q_j_uni) / 2 - matches.append((i, j, q_match)) - - return matches - - def flag_outliers_2d(self, multiplier=2, q_weight=1000, theta_weight=111): - costs = [] - for p in self.sublattice_indexed: - costs.append(self.sub_basis.compute_pair_cost( - self.sub_basis.vectors, p, q_weight, theta_weight)) - for p, c in zip(self.sublattice_indexed, costs): - p.is_outlier = c > multiplier*np.median(costs) - def plot_costs_2d(self, q_weight=1000, theta_weight=111): - costs = [] - for p in self.sublattice_indexed: - costs.append(self.sub_basis.compute_pair_cost( - self.sub_basis.vectors, p, q_weight, theta_weight, use_outliers=True)) - costs_inlier = [c for c,p in zip(costs, self.sublattice_indexed) if not p.is_outlier] - costs_outlier = [c for c,p in zip(costs, self.sublattice_indexed) if p.is_outlier] - bins = np.linspace(0, max(costs), 50) - plt.hist(costs_inlier, bins=bins) - plt.hist(costs_outlier, bins=bins, color='red') - plt.show() - - - def reprocess_all_pairs(self) -> None: - """Clear all categorizations and reprocess all pairs except the last.""" - stored_pairs = self.all_pairs # [:-1] - self.sublattice_indexed = [] - self.sublattice_1vector = [] - self.unindexed = [] - for old_pair in stored_pairs: - result, status = self.sub_basis.match(old_pair) - self._store_result(result, status) - - def _store_result(self, result: SpotPair, status: str) -> None: - """Store a processed pair in appropriate category.""" - if status == 'indexed_2d': - self.sublattice_indexed.append(result) - elif status == 'indexed_3d': - self.indexed.append(result) - - elif status == 'one_vector': - self.sublattice_1vector.append(result) - else: # unindexed - self.unindexed.append(result) - - def generate_3d_basis_indices(self, pair1_idx, pair2_idx, delta_theta_1=None, delta_theta_2=None, inv=False): - pair1 = self.sublattice_1vector[pair1_idx] - pair2 = self.sublattice_1vector[pair2_idx] - return generate_3d_basis_pairs(pair1, pair2, delta_theta_1, delta_theta_2, inv) - - def generate_3d_basis_pairs(self, pair1, pair2, delta_theta_1=None, delta_theta_2=None, inv=False): - """Generate a 3D basis from two matching one-vector pairs.""" - - # Get indexed vectors from sublattice (2D) - v1 = np.hstack((pair1.hkl1 @ self.sub_basis.vectors, 0)) - v2 = np.hstack((pair2.hkl1 @ self.sub_basis.vectors, 0)) - - # Get common q value and both angles - q = pair1.q2 - theta1 = pair1.theta_rad - theta2 = pair2.theta_rad - - parity = -1 if inv else 1 - if delta_theta_1 is not None: - theta1 += np.radians(delta_theta_1) - if delta_theta_2 is not None: - theta2 += np.radians(delta_theta_2) - # Find possible positions for third vector (returns 3D vectors) - try: - v3_candidates = [third_vector(q, v1, parity*v2, theta1, theta2)] - except ValueError: - return None - - # Choose best third vector - v3 = find_best_third_vector(v3_candidates, self.sub_basis.vectors) - - # Create full 3D basis by extending 2D vectors - v1v2 = np.hstack((self.sub_basis.vectors, np.zeros((2,1)))) - basis_vectors = np.vstack((v1v2, v3)) - - result = Basis3d.from_vectors( - vectors=basis_vectors, - qmax=self.sub_basis.qmax, - q_tol=self.sub_basis.q_tolerance, - theta_tol_deg=np.degrees(self.sub_basis.theta_tolerance) - ) - return result - - def print_status(self, verbose=False) -> None: - """Print current status of reconstruction.""" - print(f"\nCurrent minimum sublattice: {self.sub_basis}") - if hasattr(self, 'full_basis'): - print(f"Current full lattice:\n{self.full_basis}") - else: - print("No full lattice determined yet.") - - print("\n---") - if not verbose: return - print(self.summarize_half_indexed_pairs()) - return - - # If we have matching pairs, show possible 3D cells - matches = self.find_matching_pairs() - if matches: - print("\nLattice completion possibilities:") - print("entry | i1 | i2 | cell | vol (ų) | % idx") - for i, (i1, i2, q) in enumerate(matches): - try: - basis3d = self.generate_3d_basis_indices(i1, i2) - # Calculate metrics - volume = np.abs(np.linalg.det(basis3d.vectors)) - n_indexed = len([p for p in self.all_pairs - if basis3d.match(p)[1] == 'indexed_3d']) - pct_indexed = 100 * n_indexed / len(self.all_pairs) - - print(f"{i}) {i1:2d} {i2:2d} {basis3d} {volume:.1f} {pct_indexed:.1f}") - except ValueError as e: - # Skip if no solution exists - #print(f"{i}) {i1:2d} {i2:2d} No valid solution") - pass - -# Grid search stuff, temporary - -# Define grid search parameters -def grid_search_3d_basis(recon, pair1_idx, pair2_idx, - delta_range=(-2.0, 2.0), steps=21, inv=False, - verbose=True): - # Create grid of delta theta values - delta_values = np.linspace(delta_range[0], delta_range[1], steps) - grid_shape = (steps, steps) - fom = np.zeros(grid_shape) - idx = np.zeros(grid_shape) - - # Total number of pairs - total_pairs = len(recon.all_pairs) - - if verbose: - display1_fn = display2_fn = tqdm - else: - display1_fn = lambda x: x - display2_fn = lambda x, **_: x - # Grid search - for i, delta1 in display1_fn(enumerate(delta_values)): - for j, delta2 in display2_fn(enumerate(delta_values), leave=False): - try: - # Generate 3D basis with current deltas - basis3d = recon.generate_3d_basis_pairs( - pair1_idx, pair2_idx, - delta_theta_1=delta1, - delta_theta_2=delta2, - inv=inv - ) - - # Count indexed pairs - hits = 0 - vol = basis3d.volume() - if vol > .0001: - for p in recon.all_pairs: - if basis3d.match(p)[1] != 'unindexed': - hits += 1 - - # Compute percentage - fom[i, j] = 100*hits / total_pairs * vol**(1/2) - idx[i, j] = 100*hits / total_pairs - - except Exception as e: - # Failed to generate basis (e.g., no valid solution) - fom[i, j] = 0 - idx[i, j] = 0 - - return delta_values, fom, idx - -# Run grid search for both inv=False and inv=True -def run_both_grid_searches(recon, pair1_idx, pair2_idx, - delta_range=(-2.0, 2.0), steps=21, verbose=True): - # print(f"Running grid search for pair indices {pair1_idx} and {pair2_idx}...") - # print(f"Delta range: {delta_range}, steps: {steps}") - - delta_values, fom_normal, pct_normal = grid_search_3d_basis( - recon, pair1_idx, pair2_idx, delta_range, steps, inv=False, - verbose=verbose - ) - - delta_values, fom_inv, pct_inv = grid_search_3d_basis( - recon, pair1_idx, pair2_idx, delta_range, steps, inv=True, - verbose=verbose - ) - - return delta_values, (fom_normal, pct_normal), (fom_inv, pct_inv) - -# Plot the results as heatmaps -def plot_grid_search_results(delta_values, results_normal, results_inv): - fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(12, 5)) - - # Custom colormap (white to blue) - colors = [(1, 1, 1), (0, 0, 1)] - cmap = LinearSegmentedColormap.from_list('WhiteToBlue', colors) - - # Determine shared vmax for consistent coloring - vmax = max(np.max(results_normal), np.max(results_inv)) - - # Plot normal results - im1 = ax1.imshow(results_normal, extent=[delta_values[0], delta_values[-1], - delta_values[0], delta_values[-1]], - origin='lower', cmap=cmap, vmin=0, vmax=vmax) - ax1.set_title('Normal (inv=False)') - ax1.set_xlabel('Delta theta 2 (degrees)') - ax1.set_ylabel('Delta theta 1 (degrees)') - - # Plot inverted results - im2 = ax2.imshow(results_inv, extent=[delta_values[0], delta_values[-1], - delta_values[0], delta_values[-1]], - origin='lower', cmap=cmap, vmin=0, vmax=vmax) - ax2.set_title('Inverted (inv=True)') - ax2.set_xlabel('Delta theta 2 (degrees)') - -# # Add colorbar -# cbar = fig.colorbar(im1, ax=[ax1, ax2], orientation='vertical', shrink=0.8) -# cbar.set_label('Indexing percentage (%)') - - # Add max value annotations - max_normal = np.max(results_normal) - max_normal_idx = np.unravel_index(np.argmax(results_normal), results_normal.shape) - delta1_normal = delta_values[max_normal_idx[0]] - delta2_normal = delta_values[max_normal_idx[1]] - - max_inv = np.max(results_inv) - max_inv_idx = np.unravel_index(np.argmax(results_inv), results_inv.shape) - delta1_inv = delta_values[max_inv_idx[0]] - delta2_inv = delta_values[max_inv_idx[1]] - - ax1.plot(delta2_normal, delta1_normal, 'r+', markersize=10) -# ax1.text(delta2_normal, delta1_normal, f' {max_normal:.1f}%', color='red') -# - ax2.plot(delta2_inv, delta1_inv, 'r+', markersize=10) -# ax2.text(delta2_inv, delta1_inv, f' {max_inv:.1f}%', color='red') - - plt.tight_layout() - return fig - -# Function to run everything and return the best parameters -def find_best_3d_basis(recon, pair1_idx, pair2_idx, - delta_range=(-2.0, 2.0), steps=21, plot=True, - verbose=True): - """pair1_idx and pair2_idx are actually pairs""" - # Run grid searches - delta_values, results_normal, results_inv = run_both_grid_searches( - recon, pair1_idx, pair2_idx, delta_range, steps, verbose=verbose - ) - fom_normal, pct_normal = results_normal - fom_inv, pct_inv = results_inv - - # Plot results - if plot: - fig = plot_grid_search_results(delta_values, fom_normal, fom_inv) - else: - fig = None - - # Find best parameters - max_normal = np.max(fom_normal) - max_normal_idx = np.unravel_index(np.argmax(fom_normal), fom_normal.shape) - max_normal_pct = pct_normal[max_normal_idx] - delta1_normal = delta_values[max_normal_idx[0]] - delta2_normal = delta_values[max_normal_idx[1]] - - max_inv = np.max(fom_inv) - max_inv_idx = np.unravel_index(np.argmax(fom_inv), fom_inv.shape) - max_inv_pct = pct_inv[max_inv_idx] - delta1_inv = delta_values[max_inv_idx[0]] - delta2_inv = delta_values[max_inv_idx[1]] - - # Choose best overall parameters - if max_normal >= max_inv: - best_params = { - 'delta_theta_1': delta1_normal, - 'delta_theta_2': delta2_normal, - 'inv': False, - 'fom': max_normal, - 'pct': max_normal_pct - } - else: - best_params = { - 'delta_theta_1': delta1_inv, - 'delta_theta_2': delta2_inv, - 'inv': True, - 'fom': max_inv, - 'pct': max_inv_pct - } - - # Generate the best basis - best_basis = recon.generate_3d_basis_pairs( - pair1_idx, pair2_idx, - delta_theta_1=best_params['delta_theta_1'], - delta_theta_2=best_params['delta_theta_2'], - inv=best_params['inv'] - ) - - idx_pct = best_params['fom']/best_basis.volume()**(1/2) - if verbose: - print("\nBest parameters:") - print(f"delta_theta_1 = {best_params['delta_theta_1']:.2f} degrees") - print(f"delta_theta_2 = {best_params['delta_theta_2']:.2f} degrees") - print(f"inv = {best_params['inv']}") - print(f"Indexing percentage = {idx_pct:.1f}%") - - - return best_basis, best_params, fig - - -def sb1_callback(triplet, recon): - """ - Analyze a selected triplet and return q-values to display. - """ - RANGE=2 - PTS=11 - - q1, q2, theta_degrees = triplet - theta_rad = np.radians(theta_degrees) -# deltas = np.linspace(-RANGE,RANGE,PTS) -# candidates = [] -# for d in deltas: -# test_triplet = [q1, q2, theta_rad+d*np.pi/180] -# b = Basis.from_params(*test_triplet) -# candidates.append((b, b.fom_2d(recon.all_pairs))) -# candidates.sort(key=lambda x:x[1], reverse=True) -# sb1 = candidates[0][0] - - - sb1 = Basis.from_params(q1, q2, theta_rad) - sb1.match_pairs(recon.all_pairs) - - # Refine the basis - for _ in range(5): - sb1.reindex_pairs() - sb1.flag_outliers() - sb1.refine() - print(f'{sb1}: {sb1.index_percent()}') - - # Compute q-values (norms of the basis points) - qvals = np.linalg.norm(sb1.points, axis=1) - - points_1 = np.vstack((sb1.q1_values, sb1.q2_values, sb1.theta_values*180/np.pi)).T - points_2 = np.vstack((sb1.q2_values, sb1.q1_values, sb1.theta_values*180/np.pi)).T - points = np.vstack((points_1, points_2)) - - # Filter to reasonable range for display - qvals = qvals[(qvals > 0.05) & (qvals < 0.9)] - - return qvals, points - - -# Pre-compute grid once globally -_I_VALS = np.array([-2, -2, -2, -2, -2, -1, -1, -1, -1, -1, 0, 0, 0, 0, - 1, 1, 1, 1, 1, 2, 2, 2, 2, 2]) -_J_VALS = np.array([-2, -1, 0, 1, 2, -2, -1, 0, 1, 2, -2, -1, 1, 2, - -2, -1, 0, 1, 2, -2, -1, 0, 1, 2]) -@njit -def reduce_2d_cell(a, b, gamma_deg): - """ - Find the reduced setting of a 2D unit cell (Numba JIT version). - - Parameters: - a, b: lengths of cell vectors - gamma_rad: angle between vectors (in radians) - - Returns: - (a_red, b_red, gamma_red): reduced cell parameters - """ - gamma_rad = np.radians(gamma_deg) - # Convert to Cartesian coordinates - cos_gamma = np.cos(gamma_rad) - sin_gamma = np.sin(gamma_rad) - v1 = np.array([a, 0.0]) - v2 = np.array([b * cos_gamma, b * sin_gamma]) - - for _ in range(100): - # Calculate all lattice vectors - n_vecs = len(_I_VALS) - vectors = np.zeros((n_vecs, 2)) - for i in range(n_vecs): - vectors[i, 0] = _I_VALS[i] * v1[0] + _J_VALS[i] * v2[0] - vectors[i, 1] = _I_VALS[i] * v1[1] + _J_VALS[i] * v2[1] - - # Calculate lengths - lengths = np.zeros(n_vecs) - for i in range(n_vecs): - lengths[i] = np.sqrt(vectors[i, 0]**2 + vectors[i, 1]**2) - - # Sort by length - sorted_indices = np.argsort(lengths) - - # Find two shortest non-parallel vectors - shortest = vectors[sorted_indices[0]] - - # Find first non-parallel vector - second_shortest = None - for i in range(1, n_vecs): - vec = vectors[sorted_indices[i]] - cross = abs(shortest[0] * vec[1] - shortest[1] * vec[0]) - if cross > 1e-10: - second_shortest = vec - break - - if second_shortest is None: - break - - # Calculate norms - new_a_sq = shortest[0]**2 + shortest[1]**2 - new_b_sq = second_shortest[0]**2 + second_shortest[1]**2 - new_a = np.sqrt(new_a_sq) - new_b = np.sqrt(new_b_sq) - - # Ensure a <= b - if new_a > new_b: - shortest, second_shortest = second_shortest.copy(), shortest.copy() - new_a, new_b = new_b, new_a - new_a_sq, new_b_sq = new_b_sq, new_a_sq - - # Check if better - v1_len_sq = v1[0]**2 + v1[1]**2 - v2_len_sq = v2[0]**2 + v2[1]**2 - - if new_a_sq < v1_len_sq - 1e-10 or ( - abs(new_a_sq - v1_len_sq) < 1e-10 and new_b_sq < v2_len_sq - 1e-10 - ): - v1 = shortest.copy() - v2 = second_shortest.copy() - else: - break - - # Calculate reduced parameters - a_red = np.sqrt(v1[0]**2 + v1[1]**2) - b_red = np.sqrt(v2[0]**2 + v2[1]**2) - cos_gamma_red = (v1[0] * v2[0] + v1[1] * v2[1]) / (a_red * b_red) - gamma_red = np.arccos(cos_gamma_red) - - # Ensure conventional choice - if gamma_red > np.pi / 2: - gamma_red = np.pi - gamma_red - - return a_red, b_red, np.degrees(gamma_red) - - - -def run(): - QMIN=.05 - QMAX=.45 - recon = LatticeReconstruction(qmax=QMAX) - - fn1 = sys.argv[1] - data = np.load(fn1)['triplets'][:,1:4] - data[:,0] = 1/data[:,0] - data[:,1] = 1/data[:,1] - data2 = np.vstack((data[:,1], data[:,0], data[:,2])).transpose() - data_orig = data - data = np.vstack((data, data2)) - mask = (data[:, 0] >= QMIN) & (data[:, 0] <= QMAX) & \ - (data[:, 1] >= QMIN) & (data[:, 1] <= QMAX) - data = data[mask] - #data_red = np.array([reduce_2d_cell(*item) for item in tqdm(data_orig)]) - #data_red2 = np.vstack((data_red[:,1], data_red[:,0], data_red[:,2])).transpose() - #data_red_all = np.vstack((data_red, data_red2)) - - - if len(sys.argv)==3 and False: - for l in open(sys.argv[2]): - recon.store_triplet(*[float(x) for x in l.split()]) - else: - cl_auto = ManualClusterer(data, n_maxima=500, qmin=QMIN, qmax=QMAX) - for val in cl_auto.kde_maxima: - recon.store_triplet(*val) - - assert sys.argv[2] in ['auto', 'manual'] - # Manual select 2d sub bases - if sys.argv[2] == 'manual': -# cl1 = ManualClusterer(data_red_all, n_maxima=0, sb1_callback=sb1_callback, recon=recon, qmin=QMIN, qmax=QMAX) - cl1 = ManualClusterer(data, n_maxima=0, sb1_callback=sb1_callback, recon=recon, qmin=QMIN, qmax=QMAX) - title = "select first sub-basis" - triplets = cl1.select_triplets(title) - # Lattice doubling test case: 0.153379 0.125106 40.434691 - triplets = triplets[-1:] # keep the last one - triplets[0][2] = np.radians(triplets[0][2]) - sb1 = Basis.from_params(*triplets[0]) - - else: - # Auto select 2d sub basis - recon.generate_2d_bases(scan_range=0,scan_pts=1) - print('fom1') - basis_fom = [ - (b, b.fom_1d(recon.all_pairs) ) - for b in recon.basis_candidates_2d - ] - basis_fom.sort(key=lambda x:x[1], reverse=True) - top_1d = basis_fom #[:500] - print('fom2') - basis_fom1_fom2 = [ - (b, f, b.fom_2d(recon.all_pairs)) - for b, f in top_1d - ] - basis_fom1_fom2.sort(key=lambda x:x[2], reverse=True) - for i, b in enumerate(basis_fom1_fom2[:20]): - print(i, b[0], round(b[1], 1), round(b[2], 1)) - i_sb1 = int(input('Sub-basis 1? [0]') or 0) - sb1 = basis_fom1_fom2[i_sb1] - sb1 = sb1[0] - - - sb1.match_pairs(recon.all_pairs) - for _ in range(5): - sb1.reindex_pairs() - sb1.flag_outliers() - sb1.refine() - #sb1.plot_costs() - - print(f'sb1 index percent: {sb1.index_percent()}') - # Test doubled/tripled sub-bases - i_cell = None - while i_cell != 0: - doubled_cells = sb1.doubled_cells() - all_cells = [sb1] + doubled_cells - print('Cell doubling selection:') - print('i \t%idx') - for i, c in enumerate(all_cells): - c.match_pairs(recon.all_pairs) - print(i, '\t', c.index_percent(), '\t', str(c)) - i_cell = int(input('Cell? [0]') or 0) - sb1 = all_cells[i_cell] - for _ in range(3): - sb1.reindex_pairs() - sb1.flag_outliers() - sb1.refine() - - - # Manual lattice expansion - if sys.argv[2]=='manual': - print('manual selection start') - qvals_1 = np.linalg.norm(sb1.points, axis=1) - - cl2 = ManualClusterer(data, n_maxima=0, qvals_1=qvals_1, qmin=QMIN, qmax=QMAX) - triplets = cl2.select_triplets(title="Lattice expansion") - assert len(triplets) == 2 - pairs = [SpotPair(q1,q2,np.radians(th)) for q1,q2,th in triplets] - matches = [sb1.match(p) for p in pairs] - indexed_pairs = [] - for m in matches: - assert m[1] == 'one_vector' - indexed_pairs.append(m[0]) - recon.sub_basis = sb1 - best_basis, best_params, fig = find_best_3d_basis(recon, *indexed_pairs, delta_range=(-2,2), steps=11) - plt.show() - - # Auto lattice expansion - else: - recon.sub_basis = sb1 - one_vec = [x for x in sb1.all_pairs if x[2]=='one_vector'] - one_vec.sort(key=lambda x:x[1].q2) - to_try = [] - for i in range(len(one_vec)-1): - pair1 = one_vec[i][1] - pair2 = one_vec[i+1][1] - if np.abs(pair1.q2-pair2.q2)<.001 and np.abs(pair1.q1-pair2.q1)>.005: - to_try.append((pair1, pair2)) - results = [] - for i, pair in enumerate(to_try): - try: - basis, params, _ = find_best_3d_basis( - recon, *pair, delta_range=(0,0), steps=1, plot=False, - verbose=False - ) - results.append((i, basis, params['fom'], params['pct'])) - except AttributeError: - pass - except Exception as e: - raise - print(i) - results.sort(key=lambda x:x[2], reverse=True) - for i_results, (i, basis, fom, pct) in enumerate(results[:30]): - print(i_results, '\t', pct, '\t', basis) - i_results_best = int(input('lattice: [0] ') or 0) - i_best = results[i_results_best][0] - best_basis, best_params, fig = find_best_3d_basis(recon, *to_try[i_best], delta_range=(-.5,.5), steps=11, plot=False) - plt.show() - - best_basis.match_pairs(recon.all_pairs) - for _ in range(6): - best_basis.reindex_pairs() - best_basis.flag_outliers() - best_basis.refine() - #best_basis.plot_costs() - # Test multiples of the chosen cell - i_cell = None - while i_cell != 0: - doubled_cells = best_basis.doubled_cells() - all_cells = [best_basis] + doubled_cells - print('Cell doubling selection:') - print('i \t%idx\tCell') - for i, c in enumerate(all_cells): - c.match_pairs(recon.all_pairs) - print(i, '\t', c.index_percent(), '\t', str(c)) - i_cell = int(input('Cell? [0]') or 0) - best_basis = all_cells[i_cell] - for _ in range(3): - best_basis.reindex_pairs() - best_basis.flag_outliers() - best_basis.refine() - #best_basis.plot_costs_components() - - cell_vals = list(best_basis.compute_direct_cell_params(best_basis.vectors).values()) - uc = uctbx.unit_cell(cell_vals) - cs = crystal.symmetry(unit_cell=uc, space_group='P1') - subgroups = metric_subgroups(cs, max_delta=3) - subsyms = [x['best_subsym'] for x in subgroups.result_groups] - - symmetrized_bases = [] - for i, subsym in enumerate(subsyms): - print(f"\n======= Test symmetry {i}/{len(subsyms)} =======") - constr_basis = Basis.from_crystal_symmetry(subsym) - constr_basis.match_pairs(recon.all_pairs) - for _ in range(3): - constr_basis.reindex_pairs() - constr_basis.flag_outliers() - constr_basis.refine(constrained=True) - symmetrized_bases.append((constr_basis, constr_basis.index_percent())) - - for i, b in enumerate(symmetrized_bases): print(str(i) + '\t' + str(b[0]) + ' ' + str(b[1])) - i_final = int(input('Choice: [0]') or 0) - final_basis = symmetrized_bases[i_final][0] - final_points_1 = np.vstack((final_basis.q1_values, final_basis.q2_values, final_basis.theta_values*180/np.pi)).T - final_points_2 = np.vstack((final_basis.q2_values, final_basis.q1_values, final_basis.theta_values*180/np.pi)).T - final_points = np.vstack((final_points_1, final_points_2)) - final_qvals = np.linalg.norm(final_basis.points, axis=1) - -# final_points_red = np.array([reduce_2d_cell(*item) for item in tqdm(final_points_1)]) -# final_points_red2 = np.vstack((final_points_red[:,1], final_points_red[:,0], final_points_red[:,2])).transpose() -# final_points_red_all = np.vstack((final_points_red, final_points_red2)) -# -# cl_final = ManualClusterer(data_red_all, n_maxima=0, qvals_1=final_qvals, qmin=QMIN, qmax=QMAX, points=final_points_red_all) - cl_final = ManualClusterer(data, n_maxima=0, qvals_1=final_qvals, qmin=QMIN, qmax=QMAX, points=final_points) - _ = cl_final.select_triplets() - - -if __name__=="__main__": - run() - - -def junk(): - return ''' -#def find_third_vector(v1: np.ndarray, v2: np.ndarray, -# q: float, theta1: float, theta2: float) -> Tuple[np.ndarray, np.ndarray]: -# """Find a vector v3 given its length and angles with v1 and v2. -# -# Args: -# v1, v2: Two known vectors in xy-plane (2D arrays) -# q: Length of vector to find -# theta1: Angle between v1 and v3 -# theta2: Angle between v2 and v3 -# -# Returns: -# Two possible 3D positions for v3 (above/below v1-v2 plane) -# """ -# # Convert 2D vectors to 3D -# v1_3d = np.array([v1[0], v1[1], 0.0]) -# v2_3d = np.array([v2[0], v2[1], 0.0]) -# -# # Normalize -# v1_unit = v1_3d / np.linalg.norm(v1_3d) -# v2_unit = v2_3d / np.linalg.norm(v2_3d) -# -# # Get normal to v1-v2 plane (will be along z-axis) -# n = np.cross(v1_unit, v2_unit) -# n = n / np.linalg.norm(n) # should be [0, 0, ±1] -# -# # Solve for components -# v1v2 = np.dot(v1_unit, v2_unit) -# A = np.array([[1, v1v2], [v1v2, 1]]) -# b = q * np.array([np.cos(theta1), np.cos(theta2)]) -# a, b = np.linalg.solve(A, b) -# -# # Find c from length condition -# c_sq = q*q - (a*a + b*b + 2*a*b*v1v2) -# if c_sq < 0: -# raise ValueError("No solution exists for these constraints") -# c = np.sqrt(c_sq) -# -# # Return both possible 3D positions -# v3_plus = a*v1_unit + b*v2_unit + c*n -# v3_minus = a*v1_unit + b*v2_unit - c*n -# -# return v3_plus, v3_minus - -#class IndexedSpotPair2d(SpotPair): -# def initialize_orientation(self, basis_vectors: np.ndarray) -> None: -# """Initialize orientation by computing angle between calculated and observed vectors.""" -# # Get calculated vectors -# q1_calc = self.indices1[0]*basis_vectors[0] + self.indices1[1]*basis_vectors[1] -# q2_calc = self.indices2[0]*basis_vectors[0] + self.indices2[1]*basis_vectors[1] -# -# # Compute angle of q1_calc from x-axis -# phi_calc = np.arctan2(q1_calc[1], q1_calc[0]) -# -# # Our observed q1 is along x-axis, so this is our basic rotation -# self.init_phi = phi_calc -# -# # Check if we need to flip orientation by comparing second vector -# q2_obs = self.q2 * np.array([np.cos(self.theta), np.sin(self.theta)]) -# R = np.array([[np.cos(self.init_phi), -np.sin(self.init_phi)], -# [np.sin(self.init_phi), np.cos(self.init_phi)]]) -# q2_obs_rot = R @ q2_obs -# -# # If distance is large, try flipping -# if np.linalg.norm(q2_obs_rot - q2_calc) > np.linalg.norm(q2_obs_rot + q2_calc): -# self.init_phi += np.pi -# -# # Optional: small local optimization to refine this initial guess -# result = minimize_scalar( -# lambda dphi: self.compute_cost(basis_vectors, dphi), -# bounds=(-0.1, 0.1), # small range around initial guess -# method='bounded' -# ) -# self.init_phi += result.x -# -# def compute_cost(self, basis_vectors: np.ndarray, delta_phi: float) -> float: -# """Compute distance between rotated observed and calculated positions.""" -# # Compute vectors from indices (fixed) -# q1_calc = self.indices1[0]*basis_vectors[0] + self.indices1[1]*basis_vectors[1] -# q2_calc = self.indices2[0]*basis_vectors[0] + self.indices2[1]*basis_vectors[1] -# -# # Create observed vectors in standard orientation (q1 along x) -# q1_obs = np.array([self.q1, 0.0]) -# q2_obs = self.q2 * np.array([np.cos(self.theta), np.sin(self.theta)]) -# -# # Total rotation to apply to observed vectors -# total_phi = self.init_phi + delta_phi -# R = np.array([[np.cos(total_phi), -np.sin(total_phi)], -# [np.sin(total_phi), np.cos(total_phi)]]) -# -# # Rotate observed vectors to match calculated -# q1_obs_rot = R @ q1_obs -# q2_obs_rot = R @ q2_obs -# -# return (np.linalg.norm(q1_obs_rot - q1_calc) + -# np.linalg.norm(q2_obs_rot - q2_calc)) - - # Methods from Basis2d to support the unused merge_bases approach -# def basis_in_3d(self, plus_x): -# """Make a 3x2 matrix that multiplies a point in this basis to give its -# 3d coordinates. Initially the basis should lie in the xy plane with -# the plus_x vector aligned along 1,0,0. -# -# Args: -# plus_x: A 2-element array specifying which vector in the 2D basis -# should be aligned with the positive x-axis in 3D. -# For example, [1,0] means align the first basis vector. -# -# Returns: -# A 3x2 matrix that transforms 2D basis coordinates to 3D coordinates. -# """ -# # Calculate which vector in the original basis should align with x-axis -# vector_to_align = plus_x[0] * self.vectors[0] + plus_x[1] * self.vectors[1] -# -# # Calculate the angle to rotate this vector to align with [1,0] -# norm = np.linalg.norm(vector_to_align) -# cos_angle = vector_to_align[0] / norm -# sin_angle = vector_to_align[1] / norm -# -# # Create the rotation matrix (clockwise rotation) -# # This will rotate the vector_to_align to the positive x-axis -# rotation = np.array([ -# [cos_angle, sin_angle], -# [-sin_angle, cos_angle] -# ]) -# -# # Apply rotation to both basis vectors -# rotated_vectors = rotation @ self.vectors.T -# -# # Create the 3x2 transformation matrix -# # The first two rows contain the rotated vectors -# # The third row contains zeros (since we're in the xy plane) -# result = np.zeros((3, 2)) -# result[0, :] = rotated_vectors[0, :] # x components -# result[1, :] = rotated_vectors[1, :] # y components -# -# return result -# -# def vec3d_aligned_rotated(self, point, plusx, rotx_deg): -# """Construct the aligned basis from above and compute 3d coordinates for the given point. -# Apply a rotation around the x-axis and return the transformed coordinates. -# -# Args: -# point: A 2-element array representing a point in the 2D basis -# plusx: Specifies which vector should align with the x-axis -# rotx_deg: Rotation angle around the x-axis in degrees -# -# Returns: -# A 3D point after transformation and rotation -# """ -# # Step 1: Get the 3D basis that aligns plusx with the x-axis -# basis_3d = self.basis_in_3d(plusx) -# -# # Step 2: Transform the 2D point to 3D -# point_3d = basis_3d @ np.array(point) -# -# # Step 3: Create the rotation matrix around the x-axis -# rotx_rad = np.radians(rotx_deg) -# cos_rx = np.cos(rotx_rad) -# sin_rx = np.sin(rotx_rad) -# -# # Rotation matrix around x-axis -# # [1 0 0 ] -# # [0 cos(θ) -sin(θ)] -# # [0 sin(θ) cos(θ)] -# rot_x = np.array([ -# [1, 0, 0], -# [0, cos_rx, -sin_rx], -# [0, sin_rx, cos_rx] -# ]) -# -# # Step 4: Apply the rotation -# rotated_point = rot_x @ point_3d -# -# return rotated_point -#@dataclass -#class GeneratedSpotPair(): -# hkl1: np.ndarray -# hkl2: np.ndarray -# q1: np.float64 -# q2: np.float64 -# theta: np.float64 - -# Unused functions -#def common_axis(b1, b2): -# """For two Basis2d objects, find the closest matching q-value and return -# the indices in each basis that give the corresponding value. -# """ -# qvals_1 = np.linalg.norm(b1.points, axis=1) -# qvals_2 = np.linalg.norm(b2.points, axis=1) -# delta_best = 999 -# i1_best = -1 -# i2_best = -1 -# for i2, val in enumerate(qvals_2): -# deltas = np.abs(qvals_1-val) -# i1 = np.argmin(deltas) -# if deltas[i1] < delta_best: -# i1_best = i1 -# i2_best = i2 -# delta_best = deltas[i1] -# hk1_best = b1.point_indices[i1_best] -# hk2_best = b2.point_indices[i2_best] -# print('delta_best: ', delta_best) -# print('hk1_best: ', hk1_best, qvals_1[i1_best]) -# print('hk2_best: ', hk2_best, qvals_2[i2_best]) -# return hk1_best, hk2_best -# -#def merge_bases(b1, b2, common_axes, angle_deg): -# """Merge two 2D bases into a single 3D basis. -# -# Args: -# b1: First 2D basis -# b2: Second 2D basis -# common_axes: Two 2-tuples ((h1,k1), (h2,k2)) where (h1,k1) in the first basis -# is equivalent to (h2,k2) in the second basis -# angle_deg: Rotation angle in degrees around the common axis -# -# Returns: -# A Basis3d object representing the merged 3D basis -# """ -# -# # Extract common axis vectors from the tuples -# common_axis_b1, common_axis_b2 = common_axes -# -# # Convert the first basis to 3D, aligning common axis with x-axis -# basis1_3d = b1.basis_in_3d(common_axis_b1) -# -# # For the second basis, align with x-axis -# basis2_3d_aligned = b2.basis_in_3d(common_axis_b2) -# -# # Create the full 3D basis vectors -# # First vector: common axis (aligned with x-axis) -# v1 = np.array([basis1_3d[0, 0], 0, 0]) -# -# # Second vector: from first basis, already aligned -# v2 = np.array([basis1_3d[0, 1], basis1_3d[1, 1], 0]) -# -# # Third vector: from second basis, rotated around x-axis -# # Get a vector from the second basis that's not aligned with x-axis -# # (i.e., the second column of basis2_3d_aligned) -# v3_pre_rotation = np.array([basis2_3d_aligned[0, 1], basis2_3d_aligned[1, 1], 0]) -# -# # Apply rotation around x-axis -# rot_rad = np.radians(angle_deg) -# cos_rx = np.cos(rot_rad) -# sin_rx = np.sin(rot_rad) -# -# rot_x = np.array([ -# [1, 0, 0], -# [0, cos_rx, -sin_rx], -# [0, sin_rx, cos_rx] -# ]) -# -# v3 = rot_x @ v3_pre_rotation -# -# # Combine the vectors into a 3D basis -# vectors_3d = np.vstack([v1, v2, v3]) -# -# import IPython;IPython.embed() -# # Create and return a Basis3d object -# return Basis3d(vectors_3d, qmax=max(b1.qmax, b2.qmax), -# q_tolerance=max(b1.q_tolerance, b2.q_tolerance), -# theta_tol_degrees=max(np.degrees(b1.theta_tolerance), -# np.degrees(b2.theta_tolerance))) -# -#def merge_bases_brute(b1, b2, common_axes, angle_deg): -# """Merge two 2D bases into a single 3D basis using a brute force approach. -# -# Args: -# b1: First 2D basis -# b2: Second 2D basis -# common_axes: Two 2-tuples ((h1,k1), (h2,k2)) where (h1,k1) in the first basis -# is equivalent to (h2,k2) in the second basis -# angle_deg: Rotation angle in degrees around the common axis -# -# Returns: -# A Basis3d object representing the merged 3D basis -# """ -# -# # Generate all index pairs in a grid (-4 to 4 in each dimension) -# index_grid = list(itertools.product(range(5), range(5))) -# -# # Generate 3D points from b1 -# points_3d_b1 = [] -# for hk in index_grid: -# point_3d = b1.vec3d_aligned_rotated(hk, common_axes[0], 0) -# points_3d_b1.append(point_3d) -# -# # Generate 3D points from b2 with rotation -# points_3d_b2 = [] -# for hk in index_grid: -# point_3d = b2.vec3d_aligned_rotated(hk, common_axes[1], angle_deg) -# points_3d_b2.append(point_3d) -# -# # Generate all pairwise sums, eliminating near-duplicates -# sum_vectors = [] -# for v1 in points_3d_b1: -# for v2 in points_3d_b2: -# sum_vec = v1 + v2 -# length = np.linalg.norm(sum_vec) -# -# # Skip zero vectors -# if length < 1e-2: -# continue -# -# # Check if this is a near-duplicate of an existing vector -# is_duplicate = False -# for existing_vec, _ in sum_vectors: -# if np.linalg.norm(sum_vec - existing_vec) < 0.01: # 0.01 Å⁻¹ threshold -# is_duplicate = True -# break -# -# # If not a duplicate, add it -# if not is_duplicate: -# sum_vectors.append((sum_vec, length)) -# -# # Sort by length -# sum_vectors.sort(key=lambda x: x[1]) -# -# # Find three non-coplanar vectors -# basis_vectors = [] -# for vec, _ in sum_vectors: -# if len(basis_vectors) == 0: -# basis_vectors.append(vec) -# elif len(basis_vectors) == 1: -# # Check if not collinear -# cross_prod = np.cross(basis_vectors[0], vec) -# if np.linalg.norm(cross_prod) > 1e-6: -# basis_vectors.append(vec) -# elif len(basis_vectors) == 2: -# # Check if not coplanar -# v1, v2 = basis_vectors -# det = np.dot(np.cross(v1, v2), vec) -# if abs(det) > 1e-6: -# basis_vectors.append(vec) -# break -# -# if len(basis_vectors) < 3: -# raise ValueError("Could not find three non-coplanar vectors from the combined bases") -# -# # Stack the vectors into a 3D basis -# vectors_3d = np.vstack(basis_vectors) -# -# # Create and return a Basis3d object -# return Basis3d(vectors_3d, qmax=max(b1.qmax, b2.qmax), -# q_tolerance=max(b1.q_tolerance, b2.q_tolerance), -# theta_tol_degrees=max(np.degrees(b1.theta_tolerance), -# np.degrees(b2.theta_tolerance))) - -#class PairMatch_ab(VectorPairMatch): -# """A vector pair where the first vector is indexed in the first basis and -# the second vector is indexed in the second basis.""" -# def __init__(self, q1, q2, theta_obs, b1, b2, common_axes): -# qvals_1 = np.linalg.norm(b1.points, axis=1) -# deltas_1 = np.abs(qvals_1 - q1) -# i1 = np.argmin(deltas_1) -# hk1 = b1.point_indices[i1] -# qvals_2 = np.linalg.norm(b2.points, axis=1) -# deltas_2 = np.abs(qvals_2 - q2) -# i2 = np.argmin(deltas_2) -# hk2 = b2.point_indices[i2] -# -# self.hk1 = hk1 -# self.hk2 = hk2 -# self.b1 = b1 -# self.b2 = b2 -# self.theta_obs = theta_obs -# self.common_ax1, self.common_ax2 = common_axes -# -# #check if one hkl should be inverted -# vec1 = b1.vec3d_aligned_rotated(hk1, common_axes[0], 0) -# vec2 = b2.vec3d_aligned_rotated(hk2, common_axes[1], 0) -# if np.dot(vec1, vec2) < 0: -# self.hk2 = -1 * self.hk2 -# -# def angle_error(self, rotx_deg): -# vec1 = self.b1.vec3d_aligned_rotated( -# self.hk1, self.common_ax1, 0 -# ) -# vec2 = self.b2.vec3d_aligned_rotated( -# self.hk2, self.common_ax2, rotx_deg -# ) -# theta_calc = angle_between(vec1, vec2) -# return abs(theta_calc - self.theta_obs) -# -# pass - - -# Previous attempts from the main run method - -# assert len(triplets) == 1 -# triplets[0][2] = np.radians(triplets[0][2]) -# sb2 = Basis.from_params(*triplets[0]) -# sb2.match_pairs(recon.all_pairs) -# import IPython;IPython.embed() -# for _ in range(3): -# sb2.reindex_pairs() -# sb2.flag_outliers() -# sb2.plot_costs() -# sb2.refine() -# -# -# -# # Manual select ab pairs -# qvals_1 = np.linalg.norm(sb1.points, axis=1) -# qvals_2 = np.linalg.norm(sb2.points, axis=1) -# cl = ManualClusterer(data, n_maxima=0, qvals_1=qvals_1, qvals_2=qvals_2, qmin=.1, qmax=.5) -# triplets = cl.select_triplets("select a-b pairs") -# print(triplets) -# #triplets = [[0.15333570792827658, 0.13718166412632835, 114.6026808444415], [0.1532716587572158, 0.143893922376171, 32.211055684990406], [0.15329187415603598, 0.1836433785315679, 77.80422388061861], [0.1533567548034525, 0.25795800966128574, 50.959598431305984], [0.1667772122979944, 0.1415369213997462, 65.44425336483228], [0.1668435395895667, 0.14401664386592716, 57.66669217209992], [0.16672231651340277, 0.25795198264263536, 30.3907182392463], [0.20049139390604614, 0.1415445237130603, 25.921552407884032], [0.22693602221707057, 0.2580249296609254, 108.14834482578436], [0.27737438425400857, 0.14161107283719673, 52.546032607023136], [0.2774193570666868, 0.14393247301009263, 44.26189630562257], [0.27722723290187934, 0.28328068977243553, 52.552673251938415], [0.3069234169084582, 0.14153689240822068, 89.97024400009249], [0.30674731494727275, 0.14385171337570893, 75.9191414162387], [0.3067820484586015, 0.2578167873779439, 54.49355516464851]] -# for t in triplets: -# if t[2]>90: -# t[2] = 180-t[2] -# -# -# -# common_axes = common_axis(sb1, sb2) -# all_errors = [] -# for t in triplets: -# q1, q2, th = t -# pair = PairMatch_ab(q1, q2, th, sb1, sb2, common_axes) -# errors = [pair.angle_error(x) for x in range(361)] -# all_errors.append(errors) -# for err in all_errors: -# plt.plot(range(361), err) -# plt.show() -# -# import IPython;IPython.embed() -# print('using 2d basis: ', best_sb[0], round(best_sb[1],2), round(best_sb[2],2)) -# recon.set_sub_basis(best_sb[0]) -# import IPython;IPython.embed() -# for _ in range(2): -# recon.sub_basis.refine(recon.sublattice_indexed) -# recon.reprocess_all_pairs() -# -# recon.print_status(verbose=True) -# -# from cluster2 import ManualClusterer -# fn2 = sys.argv[2] -# data = np.load(fn2)['triplets'][:,1:4] -# data[:,0] = 1/data[:,0] -# data[:,1] = 1/data[:,1] -# data2 = np.vstack((data[:,1], data[:,0], data[:,2])).transpose() -# data = np.vstack((data, data2)) -# sublattice_qvals = np.linalg.norm(recon.sub_basis.points, axis=1) -# accept = False -# while not accept: -# cl = ManualClusterer(data, n_maxima=0, mark_qvals=sublattice_qvals) -# triplets = cl.select_triplets() -# assert len(triplets)==2 -# pairs = [SpotPair(q1,q2,np.radians(th)) for q1,q2,th in triplets] -# match_results = [recon.sub_basis.match(p) for p in pairs] -# matches_1v = [] -# for m in match_results: -# assert m[1] == 'one_vector' -# matches_1v.append(m[0]) -# best_basis, best_params, fig = find_best_3d_basis(recon, *matches_1v, delta_range=(-2.0, 2.0), steps=11) -# matches = [] -# hits = 0 -# for p in recon.all_pairs: -# result = best_basis.match(p) -# if result[1]=='indexed_3d': -# hits += 1 -# matches.append(result[0]) -# -# -# import IPython;IPython.embed() -# -# -# -# -# exit() -# #recon.read_triplet(*[float(x) for x in l.split()]) -# for _ in range(2): -# recon.sub_basis.refine(recon.sublattice_indexed) -# recon.reprocess_all_pairs() -''' From 68860c4f1f345ea40ffd45a3098896972330f99a Mon Sep 17 00:00:00 2001 From: Daniel Paley Date: Thu, 26 Mar 2026 16:51:21 -0400 Subject: [PATCH 04/15] Add n_max parameter to powder_from_spots for limiting plotted experiments This allows users to stop plotting after a specified number of experiments, useful when working with large datasets. Co-Authored-By: Claude Sonnet 4.5 --- xfel/small_cell/command_line/powder_from_spots.py | 3 +++ 1 file changed, 3 insertions(+) diff --git a/xfel/small_cell/command_line/powder_from_spots.py b/xfel/small_cell/command_line/powder_from_spots.py index 933fbe1c51a..904722017d4 100644 --- a/xfel/small_cell/command_line/powder_from_spots.py +++ b/xfel/small_cell/command_line/powder_from_spots.py @@ -95,6 +95,9 @@ .type = space_group .help = Show positions of miller indices from this unit_cell and space \ group. Not implemented. + n_max = None + .type = int + .help = Stop plotting after this many experiments. filter { enable = False .type = bool From ba640e66c1bb805d8f3c80e754b28859a28db73a Mon Sep 17 00:00:00 2001 From: Daniel Paley Date: Fri, 27 Mar 2026 11:14:47 -0400 Subject: [PATCH 05/15] powder_util: remove angle_histogram feature and fix n_max off-by-one Remove the superseded angle_histogram feature entirely: the angle() and create_pairwise_plots() top-level functions, the per-experiment and per-reflection angle accumulation code, and the post-loop histogram plot. Also remove import itertools (only used by that code) and the duplicate import numpy as np. Fix n_max off-by-one: change i>n_max to i>=n_max so that n_max=10 processes exactly 10 experiments (indices 0-9). --- xfel/small_cell/powder_util.py | 148 +-------------------------------- 1 file changed, 1 insertion(+), 147 deletions(-) diff --git a/xfel/small_cell/powder_util.py b/xfel/small_cell/powder_util.py index 3eedb6de71d..fe7adce447d 100644 --- a/xfel/small_cell/powder_util.py +++ b/xfel/small_cell/powder_util.py @@ -3,72 +3,12 @@ import numpy as np import matplotlib.pyplot as plt import matplotlib.ticker as tick -import numpy as np import copy from dials.array_family import flex from cctbx import uctbx from scitbx.math import five_number_summary from cctbx.crystal import symmetry import cctbx.miller -import itertools - -def angle(v1, v2): - """ - Compute the angle between two cartesian vectors. - - Parameters: - v1 (numpy.ndarray): The first vector. - v2 (numpy.ndarray): The second vector. - - Returns: - float: The angle between the vectors in degrees. - """ - dot_product = np.dot(v1, v2) - magnitude_v1 = np.linalg.norm(v1) - magnitude_v2 = np.linalg.norm(v2) - cos_theta = dot_product / (magnitude_v1 * magnitude_v2) - return np.degrees(np.arccos(cos_theta)) - -def create_pairwise_plots(points, labels): - # Create a figure with 3 subplots - fig, (ax1, ax2, ax3) = plt.subplots(1, 3, figsize=(15, 5)) - - # Define custom color map - color_map = { - -1: 'lightgray', # outliers in light gray - 0: 'red', - 1: 'blue', - 2: 'orange', - 3: 'green' - } - - # Convert labels to colors - colors = [color_map[label] if label in color_map else 'gray' for label in labels] - - # Define scatter plot properties - scatter_kwargs = {'c': labels, 'cmap': 'viridis', 's': 1, 'alpha': 0.6} - scatter_kwargs = {'c': colors, 's': 1, 'alpha': 0.6} - - # Plot ab - ax1.scatter(points[:, 0], points[:, 1], **scatter_kwargs) - ax1.set_xlabel('a') - ax1.set_ylabel('b') - ax1.set_title('a vs b') - - # Plot ac - ax2.scatter(points[:, 0], points[:, 2], **scatter_kwargs) - ax2.set_xlabel('a') - ax2.set_ylabel('c') - ax2.set_title('a vs c') - - # Plot bc - ax3.scatter(points[:, 1], points[:, 2], **scatter_kwargs) - ax3.set_xlabel('b') - ax3.set_ylabel('c') - ax3.set_title('b vs c') - - plt.tight_layout() - plt.show() class Spotfinder_radial_average: @@ -160,15 +100,6 @@ def calculate(self): assert compare_detector(ref_detector, expt.detector) expt.detector = detector - if params.angle_histogram.enable: - if 's1' not in refls.keys(): - refls.centroid_px_to_mm(expts) - refls.map_centroids_to_reciprocal_space(expts) - angles_12 = [] - angles_13 = [] - angles_23 = [] - detplot_counter = 0 - for i, expt in enumerate(expts): self.current_panelsums = [ np.zeros(params.n_bins) for _ in range(self.n_panels) @@ -184,7 +115,7 @@ def calculate(self): else: self.use_current_expt = True if i % 1000 == 0: print("experiment ", i) - if self.params.n_max is not None and i>self.params.n_max: + if self.params.n_max is not None and i >= self.params.n_max: break s0 = expt.beam.get_s0() sel = refls['id'] == i @@ -193,12 +124,6 @@ def calculate(self): intensities = refls_sel['intensity.sum.value'] panels = refls_sel['panel'] - if params.angle_histogram.enable: - r1max, r1min = params.angle_histogram.range1 - r2max, r2min = params.angle_histogram.range2 - r3max, r3min = params.angle_histogram.range3 - - i_r1, i_r2 ,i_r3 = [],[],[] for i_refl in range(len(refls_sel)): self.expt_count += 1 i_panel = panels[i_refl] @@ -211,61 +136,6 @@ def calculate(self): else: value = 1 res = self._process_pixel(i_panel, s0, panel, xy, value) - if params.angle_histogram.enable and r1max > res > r1min: - i_r1.append(i_refl) - if params.angle_histogram.enable and r2max > res > r2min: - i_r2.append(i_refl) - if params.angle_histogram.enable and r3max > res > r3min: - i_r3.append(i_refl) - - if params.angle_histogram.enable and i_r1 and i_r2 and i_r3: - i_all = i_r1 + i_r2 + i_r3 - a12, a13, a23 = [],[],[] - subsel_mask = flex.bool([n in i_all for n in range(len(refls_sel))]) - subsel_mask_inv = flex.bool([not x for x in subsel_mask]) - -# subsel = refls_sel.select(subsel_mask) -# subsel_inv = refls_sel.select(subsel_mask_inv) -# xyz1 = np.array(subsel['xyzobs.mm.value']) -# xyz2 = np.array(subsel_inv['xyzobs.mm.value']) -# plt.scatter(xyz1[:,0], xyz1[:,1], c='red', s=2) -# plt.scatter(xyz2[:,0], xyz2[:,1], c='blue', s=2) -# bc = expt.detector[0].get_beam_centre(expt.beam.get_s0()) -# plt.scatter(*bc, c='k', s=5) -# plt.xlim((100,240)) -# plt.ylim((100,240)) - - s0 = np.array(expt.beam.get_s0()) - for i1, i2, i3 in itertools.product(i_r1, i_r2, i_r3): - v1 = np.array(refls_sel[i1]['s1']) - s0 - v2 = np.array(refls_sel[i2]['s1']) - s0 - v3 = np.array(refls_sel[i3]['s1']) - s0 - a12.append(angle(v1, v2)) - a13.append(angle(v1, v3)) - a23.append(angle(v2, v3)) -# print('\n----------------------') -# print('Pairwise angles:') -# headers = '1,2: 59, 121', '1,3: 30, 150', '2,3: 25, 155' -# for header, vals in zip(headers, (a12, a13, a23)): -# print() -# print(header) -# for v in vals: print(round(v, 2)) -# plt.show() - angles_12.extend(a12) - angles_13.extend(a13) - angles_23.extend(a23) -# for i1, i2 in itertools.product(i_r1, i_r2): -# v1 = np.array(refls_sel[i1]['s1']) - s0 -# v2 = np.array(refls_sel[i2]['s1']) - s0 -# angles_12.append(angle(v1, v2)) -# for i1, i2 in itertools.product(i_r1, i_r3): -# v1 = np.array(refls_sel[i1]['s1']) - s0 -# v2 = np.array(refls_sel[i2]['s1']) - s0 -# angles_13.append(angle(v1, v2)) -# for i1, i2 in itertools.product(i_r2, i_r3): -# v1 = np.array(refls_sel[i1]['s1']) - s0 -# v2 = np.array(refls_sel[i2]['s1']) - s0 -# angles_23.append(angle(v1, v2)) for i in range(len(self.panelsums)): self.panelsums[i] = self.panelsums[i] + self.current_panelsums[i] @@ -285,22 +155,6 @@ def calculate(self): for i in range(len(self.panelsums)): self.antifiltered_panelsums[i] = \ self.antifiltered_panelsums[i] + self.current_panelsums[i] - if params.angle_histogram.enable: -# from sklearn.cluster import DBSCAN -# data=np.vstack((angles_12, angles_13, angles_23)).transpose() -# dbscan = DBSCAN(eps=6, min_samples=15) -# labels = dbscan.fit_predict(data) -# create_pairwise_plots(data, labels) -# fig, (ax1, ax2, ax3) = plt.subplots(1,3) -# ax1.scatter(angles_12, angles_13, s=.5) -# ax2.scatter(angles_13, angles_23, s=.5) -# ax3.scatter(angles_12, angles_23, s=.5) - fig, (ax1, ax2, ax3) = plt.subplots(3,1) - ax1.hist(angles_12, bins=180) - ax2.hist(angles_13, bins=180) - ax3.hist(angles_23, bins=180) - - plt.show() self.dvals = np.array(self.dvals) From 5c782462a53690a920770989ed33f50a1562d3c0 Mon Sep 17 00:00:00 2001 From: Daniel Paley Date: Thu, 2 Apr 2026 15:26:28 -0400 Subject: [PATCH 06/15] powder_refine_geometry: use plain str for space_group in print --- xfel/small_cell/command_line/powder_refine_geometry.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/xfel/small_cell/command_line/powder_refine_geometry.py b/xfel/small_cell/command_line/powder_refine_geometry.py index 37a51140823..26338239c55 100644 --- a/xfel/small_cell/command_line/powder_refine_geometry.py +++ b/xfel/small_cell/command_line/powder_refine_geometry.py @@ -130,7 +130,7 @@ def run(self): if params.unit_cell is not None and params.space_group is not None: print(f"Using unit_cell: {params.unit_cell}") - print(f"Using space_group: {params.space_group.symbol_and_number()}") + print(f"Using space_group: {params.space_group}") else: print(f"Reference d-spacings: {params.reference_d_spacings}") From a558f11b666b863fdbdbbc8a2e26076644a1f0b5 Mon Sep 17 00:00:00 2001 From: Daniel Paley Date: Thu, 2 Apr 2026 15:26:43 -0400 Subject: [PATCH 07/15] geometry_refiner: hierarchy access, flex.double fix, remove average_unit_cell --- xfel/small_cell/geometry_refiner.py | 30 +++++++++++++++-------------- 1 file changed, 16 insertions(+), 14 deletions(-) diff --git a/xfel/small_cell/geometry_refiner.py b/xfel/small_cell/geometry_refiner.py index b74f7fe41e6..168b52659a9 100644 --- a/xfel/small_cell/geometry_refiner.py +++ b/xfel/small_cell/geometry_refiner.py @@ -30,7 +30,7 @@ def __init__(self, experiments, reflections, params): # Compute reference d-spacings from unit_cell and space_group if provided if params.unit_cell is not None and params.space_group is not None: print(f"Computing d-spacings from unit_cell={params.unit_cell} " - f"and space_group={params.space_group.symbol_and_number()}") + f"and space_group={params.space_group}") # Get beam and detector for resolution calculation beam = experiments[0].beam @@ -71,14 +71,16 @@ def __init__(self, experiments, reflections, params): self._prepare_reflections() def _get_detector_state(self): - """Extract fast, slow, origin, and center from detector panel.""" - panel = self.detector[0] - fast = matrix.col(panel.get_fast_axis()) - slow = matrix.col(panel.get_slow_axis()) - origin = matrix.col(panel.get_origin()) + """Extract fast, slow, origin from detector hierarchy.""" + hierarchy = self.detector.hierarchy() + fast = matrix.col(hierarchy.get_local_fast_axis()) + slow = matrix.col(hierarchy.get_local_slow_axis()) + origin = matrix.col(hierarchy.get_local_origin()) normal = fast.cross(slow) - # Compute panel center + # For multipanel detectors, compute center using first panel dimensions + # (this is just for rotation center - all panels move together) + panel = self.detector[0] size = panel.get_image_size() pixel_size = panel.get_pixel_size() center = origin + (size[0]/2 * pixel_size[0]) * fast + (size[1]/2 * pixel_size[1]) * slow @@ -216,9 +218,9 @@ def apply_params(self, x): shift2 * d2 + dist * dn) - # Update detector panel - panel = self.detector[0] - panel.set_frame( + # Update detector hierarchy (moves all panels together as rigid body) + hierarchy = self.detector.hierarchy() + hierarchy.set_local_frame( d1_new.elems, d2_new.elems, new_origin.elems @@ -299,10 +301,10 @@ def get_refined_experiments(self): def report_geometry_changes(self): """Print summary of geometry changes.""" - panel = self.detector[0] - new_origin = matrix.col(panel.get_origin()) - new_fast = matrix.col(panel.get_fast_axis()) - new_slow = matrix.col(panel.get_slow_axis()) + hierarchy = self.detector.hierarchy() + new_origin = matrix.col(hierarchy.get_local_origin()) + new_fast = matrix.col(hierarchy.get_local_fast_axis()) + new_slow = matrix.col(hierarchy.get_local_slow_axis()) old_origin = self.initial_state['origin'] old_fast = self.initial_state['fast'] From 5e81ce8f131e0d03852c75b58c56d5657dc7f488 Mon Sep 17 00:00:00 2001 From: Daniel Paley Date: Thu, 2 Apr 2026 15:26:49 -0400 Subject: [PATCH 08/15] powder_util: add search_step_z, mean_squared_error_from_target, net_distance_shift --- xfel/small_cell/powder_util.py | 79 ++++++++++++++++++++++++++++++++++ 1 file changed, 79 insertions(+) diff --git a/xfel/small_cell/powder_util.py b/xfel/small_cell/powder_util.py index fe7adce447d..18343b8f4a4 100644 --- a/xfel/small_cell/powder_util.py +++ b/xfel/small_cell/powder_util.py @@ -279,6 +279,7 @@ def __init__(self, experiments, reflections, params): assert len(self.experiments.detectors())==1 self.net_origin_shift = np.array([0.,0.,0.]) + self.net_distance_shift = 0.0 self.centroid_px_mm_done = False self.px_size = self.experiments.detectors()[0][0].get_pixel_size() self.target_refl_count = 0 @@ -327,6 +328,16 @@ def width(self): print(f'width {result:.5f} from {sel.count(True)} dvals') return result + def mean_squared_error_from_target(self, d_target): + """Calculate mean squared difference between d-spacings and target d-spacing""" + if len(self.dvals) < 3: return 999 + dvals_array = flex.double(self.dvals) + differences = dvals_array - d_target + mse = flex.mean(differences * differences) + rmse = np.sqrt(mse) + print(f'RMSE from target {d_target:.5f}: {rmse:.5f} from {len(self.dvals)} dvals') + return mse + def search_step(self, step_px, nsteps=3, update=True): step_size_mm = np.array(self.px_size + (0.,)) * step_px assert nsteps%2 == 1, "nsteps should be odd" @@ -368,6 +379,74 @@ def search_step(self, step_px, nsteps=3, update=True): print(f'net shift: {self.net_origin_shift}') return width_start, width_end, self.net_origin_shift + def search_step_z(self, step_um, d_target, nsteps=3, update=True, min_refl_fraction=0.8): + """Search along detector z-axis (distance) to minimize MSE from d_target + + Args: + step_um: step size in microns + d_target: target d-spacing value + nsteps: number of steps to search (default 3: -1, 0, +1) + update: whether to update the geometry with the best result + min_refl_fraction: reject steps if reflection count drops below this fraction (default 0.8) + """ + step_size_mm = step_um / 1000.0 # convert microns to mm + assert nsteps % 2 == 1, "nsteps should be odd" + step_min = -1 * (nsteps // 2) + step_max = nsteps // 2 + 1e-6 # make the range inclusive + step_arange = np.arange(step_min, step_max) + + detector = self.experiments.detectors()[0] + hierarchy = detector.hierarchy() + fast = hierarchy.get_local_fast_axis() + slow = hierarchy.get_local_slow_axis() + origin = hierarchy.get_local_origin() + + # Compute detector normal (z-axis direction) + # Normal points from sample toward detector + normal = np.cross(fast, slow) + normal = normal / np.linalg.norm(normal) + + results = [] + self.update_dvals() + initial_count = len(self.dvals) + mse_start = self.mean_squared_error_from_target(d_target) + print(f'Z-axis search start (MSE): {mse_start:.5f}, n_refl: {initial_count}') + + for step_idx in step_arange: + step_mm = step_idx * step_size_mm + # Shift along normal direction + origin_shift = normal * step_mm + new_origin = origin + origin_shift + hierarchy.set_local_frame(fast, slow, new_origin) + self.update_dvals() + current_count = len(self.dvals) + + # Reject if reflection count dropped by more than (1 - min_refl_fraction) + if current_count < min_refl_fraction * initial_count: + result = 999 # penalty value + print(f' Step {step_mm:.6f} mm: rejected (n_refl={current_count}, {100*current_count/initial_count:.1f}%)') + else: + result = self.mean_squared_error_from_target(d_target) + results.append(result) + + mse_end = min(results) + i_best = results.index(mse_end) + distance_shift = step_arange[i_best] * step_size_mm + if update: + origin_shift = normal * distance_shift + new_origin = origin + origin_shift + self.net_origin_shift += origin_shift + self.net_distance_shift += distance_shift + else: + new_origin = origin + hierarchy.set_local_frame(fast, slow, new_origin) + self.update_dvals() + final_count = len(self.dvals) + print(f'Z step: {distance_shift:.6f} mm') + print(f'Z end (MSE): {mse_end:.5f}, n_refl: {final_count}') + print(f'net distance shift: {self.net_distance_shift:.6f} mm') + return mse_start, mse_end, self.net_distance_shift + def augment(expts, refls, d_min, d_max): """ Add pairwise 3D spot distances to the d-spacing histogram """ lab = flex.vec3_double() From 84feb91240fecb6833c8358ce98160de126084e4 Mon Sep 17 00:00:00 2001 From: Daniel Paley Date: Thu, 2 Apr 2026 15:28:14 -0400 Subject: [PATCH 09/15] powder_from_spots: add q-parameters and Z-axis refinement --- .../command_line/powder_from_spots.py | 30 +++++++++++++++++++ 1 file changed, 30 insertions(+) diff --git a/xfel/small_cell/command_line/powder_from_spots.py b/xfel/small_cell/command_line/powder_from_spots.py index 904722017d4..c23d9cd4174 100644 --- a/xfel/small_cell/command_line/powder_from_spots.py +++ b/xfel/small_cell/command_line/powder_from_spots.py @@ -142,9 +142,21 @@ .type = float d_max = None .type = float + q_min = None + .type = float + .help = Minimum q (inverse angstroms). Converted to d_max if provided. + q_max = None + .type = float + .help = Maximum q (inverse angstroms). Converted to d_min if provided. step_px = None .type = float .multiple = True + d_target = None + .type = float + .help = If set, enables detector distance refinement along z-axis + q_target = None + .type = float + .help = Target q value (inverse angstroms). Converted to d_target if provided. } plot { interactive = True @@ -188,11 +200,29 @@ def run(self): experiments = params.input.experiments[0].data reflections = params.input.reflections[0].data + if params.center_scan.q_min is not None: + params.center_scan.d_max = 1.0 / params.center_scan.q_min + if params.center_scan.q_max is not None: + params.center_scan.d_min = 1.0 / params.center_scan.q_max + if params.center_scan.q_target is not None: + params.center_scan.d_target = 1.0 / params.center_scan.q_target + if params.center_scan.d_min: assert params.center_scan.d_max cscan = Center_scan(experiments, reflections, params) + + # First XY refinement sequence for step in params.center_scan.step_px: cscan.search_step(step) + + if params.center_scan.d_target is not None: + z_steps_um = [4000, 2000, 1000, 500, 250, 125, 80, 40, 20, 10, 5] + for step_um in z_steps_um: + cscan.search_step_z(step_um, params.center_scan.d_target) + # Second XY refinement sequence + for step in params.center_scan.step_px: + cscan.search_step(step) + if params.output.geom_file is not None: experiments.as_file(params.output.geom_file) From 2986238307f160af8e24505fe5246b5bdcd79719 Mon Sep 17 00:00:00 2001 From: Daniel Paley Date: Thu, 2 Apr 2026 15:28:49 -0400 Subject: [PATCH 10/15] geometry_refiner: vectorized residuals, experiment subset, progress counter, relaxed tolerances --- xfel/small_cell/geometry_refiner.py | 56 +++++++++++++++++++++++++---- 1 file changed, 50 insertions(+), 6 deletions(-) diff --git a/xfel/small_cell/geometry_refiner.py b/xfel/small_cell/geometry_refiner.py index b74f7fe41e6..26ce5427a08 100644 --- a/xfel/small_cell/geometry_refiner.py +++ b/xfel/small_cell/geometry_refiner.py @@ -166,6 +166,32 @@ def _prepare_reflections(self): self.panels = self.refls_filtered['panel'] self.ids = self.refls_filtered['id'] + # Pre-compute unique experiment IDs and create subset for optimization + unique_exp_ids_list = sorted(set(self.refls_filtered['id'])) + print(f"Filtered reflections span {len(unique_exp_ids_list)} unique experiments " + f"out of {len(self.experiments)}") + + # Create a subset of experiments and remap reflection IDs + # This dramatically speeds up coordinate transformations + if len(unique_exp_ids_list) < len(self.experiments): + from dxtbx.model.experiment_list import ExperimentList + self.experiments_subset = ExperimentList() + old_to_new_id = {} + for new_id, old_id in enumerate(unique_exp_ids_list): + self.experiments_subset.append(self.experiments[old_id]) + old_to_new_id[old_id] = new_id + + # Remap reflection IDs to the subset + old_ids = self.refls_filtered['id'] + new_ids = flex.size_t([old_to_new_id[old_id] for old_id in old_ids]) + self.refls_filtered['id'] = new_ids + + # Use the subset for all operations + self.experiments = self.experiments_subset + print(f"Created experiment subset with {len(self.experiments)} experiments") + else: + print(f"Using all experiments (no subset needed)") + def apply_params(self, x): """Apply parameter vector to detector geometry.""" # Parse parameter vector @@ -243,17 +269,34 @@ def find_nearest_reference(self, d_obs): distances = np.abs(self.reference_d - d_obs) return self.reference_d[np.argmin(distances)] + def find_nearest_references_vectorized(self, d_obs_array): + """Find closest reference d-spacing for array of observations (vectorized).""" + # d_obs_array shape: (n_obs,) + # reference_d shape: (n_ref,) + # Compute distances matrix: (n_obs, n_ref) + distances = np.abs(d_obs_array[:, np.newaxis] - self.reference_d[np.newaxis, :]) + # Find index of minimum distance for each observation + nearest_indices = np.argmin(distances, axis=1) + return self.reference_d[nearest_indices] + def objective(self, x): """Compute sum of squared residuals.""" self.apply_params(x) dvals = self.compute_dvals() - residuals = [] - for d_obs in dvals: - d_ref = self.find_nearest_reference(d_obs) - residuals.append(d_obs - d_ref) + # Vectorized computation of residuals + dvals_np = dvals.as_numpy_array() if hasattr(dvals, 'as_numpy_array') else np.array(dvals) + d_ref = self.find_nearest_references_vectorized(dvals_np) + residuals = dvals_np - d_ref + + # Progress counter + if not hasattr(self, '_obj_eval_count'): + self._obj_eval_count = 0 + self._obj_eval_count += 1 + if self._obj_eval_count % 10 == 0: + print(f" Objective evaluation {self._obj_eval_count}, value: {np.sum(residuals ** 2):.6f}") - return np.sum(np.array(residuals) ** 2) + return np.sum(residuals ** 2) def run(self): """Run refinement and return results.""" @@ -272,7 +315,7 @@ def run(self): self.objective, x0=self.initial_params, method='Powell', - options={'maxiter': 100, 'disp': True, 'xtol': 0.001, 'ftol': 0.0001} + options={'maxiter': 50, 'disp': True, 'xtol': 0.01, 'ftol': 0.001} ) # Apply final parameters @@ -283,6 +326,7 @@ def run(self): final_obj = result.fun print(f"\nFinal objective: {final_obj:.6f}") print(f"Final RMS residual: {np.sqrt(final_obj / len(self.refls_filtered)):.6f} A") + print(f"Total objective evaluations: {self._obj_eval_count}") # Report parameter changes print("\nRefined parameters:") From b154a07ffbdb614e01e2607b407890c276ac1ca5 Mon Sep 17 00:00:00 2001 From: Daniel Paley Date: Thu, 2 Apr 2026 15:36:30 -0400 Subject: [PATCH 11/15] Add missing comments for q-parameter conversion and Z-refinement blocks --- xfel/small_cell/command_line/powder_from_spots.py | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/xfel/small_cell/command_line/powder_from_spots.py b/xfel/small_cell/command_line/powder_from_spots.py index c23d9cd4174..959a3f19fbe 100644 --- a/xfel/small_cell/command_line/powder_from_spots.py +++ b/xfel/small_cell/command_line/powder_from_spots.py @@ -200,6 +200,7 @@ def run(self): experiments = params.input.experiments[0].data reflections = params.input.reflections[0].data + # Convert q parameters to d parameters if needed (q = 1/d) if params.center_scan.q_min is not None: params.center_scan.d_max = 1.0 / params.center_scan.q_min if params.center_scan.q_max is not None: @@ -215,10 +216,13 @@ def run(self): for step in params.center_scan.step_px: cscan.search_step(step) + # Z-axis refinement if d_target is set if params.center_scan.d_target is not None: + # Default Z steps in microns z_steps_um = [4000, 2000, 1000, 500, 250, 125, 80, 40, 20, 10, 5] for step_um in z_steps_um: cscan.search_step_z(step_um, params.center_scan.d_target) + # Second XY refinement sequence for step in params.center_scan.step_px: cscan.search_step(step) From 107d3c689ba1645269d4d6fdc1cda1a77b6164ce Mon Sep 17 00:00:00 2001 From: Daniel Paley Date: Thu, 2 Apr 2026 15:55:56 -0400 Subject: [PATCH 12/15] powder_util: restore matplotlib backend assertion for interactive peak picking --- xfel/small_cell/powder_util.py | 7 +++---- 1 file changed, 3 insertions(+), 4 deletions(-) diff --git a/xfel/small_cell/powder_util.py b/xfel/small_cell/powder_util.py index 18343b8f4a4..96646a546fc 100644 --- a/xfel/small_cell/powder_util.py +++ b/xfel/small_cell/powder_util.py @@ -222,10 +222,9 @@ def plot(self): if params.plot.interactive and params.output.peak_file: backend_list = ["TkAgg","QtAgg"] - print(plt.get_backend()) -# assert (plt.get_backend() in backend_list), """Matplotlib backend not compatible with interactive peak picking. -#You can set the MPLBACKEND environment varibale to change this. -#Currently supported options: %s""" %backend_list + assert (plt.get_backend() in backend_list), """Matplotlib backend not compatible with interactive peak picking. +You can set the MPLBACKEND environment variable to change this. +Currently supported options: %s""" %backend_list #If a peak list output file is specified, do interactive peak picking: with open(params.output.peak_file, 'w') as f: vertical_line = ax.axvline(color='r', lw=0.8, ls='--', x=xvalues[1]) From 32cf2678908e2371c99f149f53331f9e5b7ef4da Mon Sep 17 00:00:00 2001 From: Daniel Paley Date: Thu, 2 Apr 2026 16:04:14 -0400 Subject: [PATCH 13/15] powder_util: revert anomalous_flag to False for powder peak overlay --- xfel/small_cell/powder_util.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/xfel/small_cell/powder_util.py b/xfel/small_cell/powder_util.py index 96646a546fc..a7977bda82e 100644 --- a/xfel/small_cell/powder_util.py +++ b/xfel/small_cell/powder_util.py @@ -207,7 +207,7 @@ def plot(self): sym = symmetry( unit_cell=params.unit_cell, space_group=params.space_group.group() ) - hkl_list = cctbx.miller.build_set(sym, True, d_min=params.d_min) + hkl_list = cctbx.miller.build_set(sym, False, d_min=params.d_min) dspacings = params.unit_cell.d(hkl_list.indices()) # for hkl, d in sorted(zip(hkl_list.indices(), dspacings), key=lambda x:x[1]): # print('{:.3f}: {}'.format(d, hkl)) From 96e3db044e0a11ed2a4a6494e8eaf973bf243420 Mon Sep 17 00:00:00 2001 From: Daniel Paley Date: Mon, 6 Apr 2026 15:49:45 -0400 Subject: [PATCH 14/15] Add xy_file_units phil param, default to d-spacing --- xfel/small_cell/command_line/powder_from_spots.py | 4 ++++ xfel/small_cell/powder_util.py | 4 +++- 2 files changed, 7 insertions(+), 1 deletion(-) diff --git a/xfel/small_cell/command_line/powder_from_spots.py b/xfel/small_cell/command_line/powder_from_spots.py index 959a3f19fbe..0f1940366ba 100644 --- a/xfel/small_cell/command_line/powder_from_spots.py +++ b/xfel/small_cell/command_line/powder_from_spots.py @@ -124,6 +124,10 @@ .type = str xy_file = None .type = str + xy_file_units = *d q + .type = choice + .help = X-axis units for xy_file output: d-spacing in Angstroms (default) \ + or inverse-d-spacing (q) in inverse Angstroms. peak_file = None .type = str .help = Optionally, specify an output file for interactive peak picking in \ diff --git a/xfel/small_cell/powder_util.py b/xfel/small_cell/powder_util.py index a7977bda82e..57530b2f604 100644 --- a/xfel/small_cell/powder_util.py +++ b/xfel/small_cell/powder_util.py @@ -198,8 +198,10 @@ def plot(self): if params.output.xy_file: with open(params.output.xy_file, 'w') as f: + use_q = getattr(params.output, 'xy_file_units', 'd') == 'q' for x,y in zip(xvalues, yvalues): - f.write("{:.6f}\t{}\n".format(x, y)) + xout = x if use_q else 1/x + f.write("{:.6f}\t{}\n".format(xout, y)) # Now plot the predicted peak positions if requested if params.unit_cell or params.space_group: From fe2a70cd45d0b9e7dec8597bd9c1a2fe94896e4d Mon Sep 17 00:00:00 2001 From: Daniel Paley Date: Mon, 6 Apr 2026 16:00:24 -0400 Subject: [PATCH 15/15] clean clutter --- xfel/small_cell/geometry_refiner.py | 1 - xfel/small_cell/small_cell.py | 2 +- 2 files changed, 1 insertion(+), 2 deletions(-) diff --git a/xfel/small_cell/geometry_refiner.py b/xfel/small_cell/geometry_refiner.py index 764460b91a1..03454dd14d7 100644 --- a/xfel/small_cell/geometry_refiner.py +++ b/xfel/small_cell/geometry_refiner.py @@ -1,6 +1,5 @@ from __future__ import division import numpy as np -import copy from dials.array_family import flex from cctbx import uctbx, miller diff --git a/xfel/small_cell/small_cell.py b/xfel/small_cell/small_cell.py index 48853e745d4..cc38186ee41 100644 --- a/xfel/small_cell/small_cell.py +++ b/xfel/small_cell/small_cell.py @@ -941,7 +941,7 @@ def small_cell_index_detail(experiments, reflections, horiz_phil, write_output = flex.random_permutation(len(all_reflections)) )[: int(len(all_reflections) * pct / 100)] - + try: lattice_results = small_cell_index_lattice_detail(experiments, reflections, horiz_phil) if not lattice_results: