From 8f3ad59c6dca1edf5f70ed0e181f54e7636babea Mon Sep 17 00:00:00 2001 From: Abe Levitan Date: Tue, 2 Dec 2025 15:27:20 +0100 Subject: [PATCH] Suggested changes: Make surface_normal work with anything castable to an array with np.asarray, change the check for length-3 to cover all cases, and add a check to explicitly fail if the scattering_mode argument is improperly set, instead of silently falling back to default --- src/cdtools/models/bragg_2d_ptycho.py | 51 ++++++++++++++------------- 1 file changed, 26 insertions(+), 25 deletions(-) diff --git a/src/cdtools/models/bragg_2d_ptycho.py b/src/cdtools/models/bragg_2d_ptycho.py index 3a981ca..b3869a9 100644 --- a/src/cdtools/models/bragg_2d_ptycho.py +++ b/src/cdtools/models/bragg_2d_ptycho.py @@ -283,33 +283,34 @@ class Bragg2DPtycho(CDIModel): # The surface normal definition is based on the following heirarchy: # manual surface_normal definition > scattering_mode # > dataset.sample_info['orientation'] > transmission geometry - - # Guard against any surface_normal entries that are not numpy vectors - if surface_normal is not None and \ - not isinstance(surface_normal, np.ndarray) or \ - (isinstance(surface_normal, np.ndarray) - and not surface_normal.shape == (3,)): - raise RuntimeError( - 'surface_normal needs to be a numpy vector with 3 elements.' + if surface_normal is not None: + surface_normal = np.asarray(surface_normal) + elif scattering_mode.strip().lower() in {'t', 'transmission'}: + surface_normal = np.array([0.,0.,1.]) + elif scattering_mode.strip().lower() in {'r', 'reflection'}: + outgoing_dir = np.cross(det_basis[:,0], det_basis[:,1]) + outgoing_dir /= np.linalg.norm(outgoing_dir) + surface_normal = outgoing_dir + np.array([0.,0.,1.]) + surface_normal /= np.linalg.norm(outgoing_dir) + elif scattering_mode is not None: + raise ValueError( + 'Scattering mode must be either "transmission" ("t"), "reflection" ("r"), or the default of None.' ) + elif hasattr(dataset, 'sample_info') and \ + dataset.sample_info is not None and \ + 'orientation' in dataset.sample_info: + # If the scattering_mode has not been defined, we grab + # this from the cxi file if its present. + surface_normal = dataset.sample_info['orientation'][2] + else: + surface_normal = np.array([0., 0., 1.]) - if surface_normal is None: - if scattering_mode in {'t', 'transmission'}: - surface_normal = np.array([0.,0.,1.]) - elif scattering_mode in {'r', 'reflection'}: - outgoing_dir = np.cross(det_basis[:,0], det_basis[:,1]) - outgoing_dir /= np.linalg.norm(outgoing_dir) - surface_normal = outgoing_dir + np.array([0.,0.,1.]) - surface_normal /= np.linalg.norm(outgoing_dir) - elif hasattr(dataset, 'sample_info') and \ - dataset.sample_info is not None and \ - 'orientation' in dataset.sample_info: - # If the scattering_mode has not been defined, we grab - # this from the cxi file if its present. - surface_normal = dataset.sample_info['orientation'][2] - else: - surface_normal = np.array([0., 0., 1.]) - + # Guard against any surface_normal entries that are not castable + # to a length-3 numpy vector, with a sensible error message + if not surface_normal.shape == (3,): + raise ValueError( + '`surface_normal` needs to be a numpy vector with 3 elements. If it was set incorrectly from dataset.sample_info, consider explicitly setting it via the `surface_normal` keyword argument.' + ) # and we use that to generate the probe basis ew_normal = np.cross(np.array(ew_basis)[:,1],