reflectorch 1.3.0__py3-none-any.whl → 1.4.0__py3-none-any.whl

reflectorch/inference/scipy_fitter.py

@@ -63,6 +63,8 @@ def refl_fit(
         bounds: np.ndarray = None,
         error_bars: np.ndarray = None,
         scale_curve_func=np.log10,
+        method: str = 'trf', #'lm', 'trf'
+        polishing_max_nfev: int = None,
         reflectivity_kwargs: dict = None,
         **kwargs
 ):
@@ -77,14 +79,23 @@ def refl_fit(
                 adjusted_bounds[1, i] += epsilon
 
         init_params = np.clip(init_params, *adjusted_bounds)
-        kwargs['bounds'] = adjusted_bounds
+        if method != 'lm':
+            kwargs['bounds'] = adjusted_bounds
 
     reflectivity_kwargs = reflectivity_kwargs or {}
     for key, value in reflectivity_kwargs.items():
         if isinstance(value, float):
             reflectivity_kwargs[key] = torch.tensor([[value]], dtype=torch.float64)
         elif isinstance(value, np.ndarray):
-            reflectivity_kwargs[key] = torch.tensor(value, dtype=torch.float32).unsqueeze(0)       
+            reflectivity_kwargs[key] = torch.tensor(value, dtype=torch.float32).unsqueeze(0)     
+
+    curve = np.clip(curve, a_min=1e-12, a_max=None)
+
+    if error_bars is not None and scale_curve_func == np.log10:
+        error_bars = np.clip(error_bars, a_min=1e-20, a_max=None)
+        scaled_error_bars = error_bars / (curve * np.log(10))
+    else:
+        scaled_error_bars = None  
 
     res = curve_fit(
         f=get_scaled_curve_func(
@@ -93,10 +104,12 @@ def refl_fit(
             reflectivity_kwargs=reflectivity_kwargs,
         ),
         xdata=q, 
-        ydata=scale_curve_func(curve),
+        ydata=scale_curve_func(curve).reshape(-1),
         p0=init_params,
-        sigma=error_bars if error_bars is not None else None,
+        sigma=scaled_error_bars,
         absolute_sigma=True,
+        method=method,
+        max_nfev=polishing_max_nfev,
         **kwargs
     )
 
@@ -184,7 +197,8 @@ def get_scaled_curve_func(
         fitted_curve = fitted_curve_tensor.squeeze().numpy()
         
         scaled_curve = scale_curve_func(fitted_curve)
-        return scaled_curve
+
+        return scaled_curve.reshape(-1)
 
     return scaled_curve_func
 

reflectorch/ml/trainers.py

@@ -24,7 +24,9 @@ class BasicBatchData:
     scaled_sigmas: Optional[torch.Tensor] = None
     scaled_q_values: Optional[torch.Tensor] = None
     scaled_denoised_curves: Optional[torch.Tensor] = None
+    key_padding_mask: Optional[torch.Tensor] = None
     scaled_conditioning_params: Optional[torch.Tensor] = None
+    unscaled_q_values: Optional[torch.Tensor] = None
 
 class RealTimeSimTrainer(Trainer):
     """Trainer with functionality to customize the sampled batch of data"""
@@ -74,6 +76,7 @@ class PointEstimatorTrainer(RealTimeSimTrainer):
         scaled_curves = batch_data['scaled_noisy_curves'].to(torch.float32)
         scaled_denoised_curves = get_scaled_or_none('curves', self.loader.curves_scaler.scale)
         scaled_q_values = get_scaled_or_none('q_values', self.loader.q_generator.scale_q) if self.train_with_q_input else None
+        key_padding_mask = batch_data.get('key_padding_mask', None)
 
         scaled_q_resolutions = get_scaled_or_none('q_resolutions', self.loader.smearing.scale_resolutions) if self.condition_on_q_resolutions else None
         conditioning_params = []
@@ -92,6 +95,8 @@ class PointEstimatorTrainer(RealTimeSimTrainer):
             scaled_q_values=scaled_q_values,
             scaled_denoised_curves=scaled_denoised_curves,
             scaled_conditioning_params=scaled_conditioning_params,
+            unscaled_q_values=batch_data['q_values'],
+            key_padding_mask=key_padding_mask,
         )
 
     def get_loss_dict(self, batch_data: BasicBatchData):
@@ -100,13 +105,17 @@ class PointEstimatorTrainer(RealTimeSimTrainer):
         scaled_curves=batch_data.scaled_curves
         scaled_bounds=batch_data.scaled_bounds
         scaled_q_values=batch_data.scaled_q_values
+        key_padding_mask=batch_data.key_padding_mask
         scaled_conditioning_params=batch_data.scaled_conditioning_params
+        unscaled_q_values=batch_data.unscaled_q_values
 
         predicted_params = self.model(
             curves = scaled_curves,
             bounds = scaled_bounds,
             q_values = scaled_q_values,
             conditioning_params = scaled_conditioning_params,
+            key_padding_mask = key_padding_mask,
+            unscaled_q_values = unscaled_q_values,
         )
 
         if not self.rescale_loss_interval_width:

reflectorch/models/__init__.py

@@ -6,6 +6,7 @@ __all__ = [
     "ConvDecoder",
     "ConvAutoencoder",
     "FnoEncoder",
+    "IntegralConvEmbedding",
     "SpectralConv1d",
     "ConvResidualNet1D",
     "ResidualMLP",

reflectorch/models/encoders/__init__.py

@@ -4,6 +4,7 @@ from reflectorch.models.encoders.conv_encoder import (
     ConvAutoencoder,
 )
 from reflectorch.models.encoders.fno import FnoEncoder, SpectralConv1d
+from reflectorch.models.encoders.integral_kernel_embedding import IntegralConvEmbedding
 from reflectorch.models.encoders.conv_res_net import ConvResidualNet1D
 
 
@@ -14,4 +15,5 @@ __all__ = [
     "ConvResidualNet1D",
     "FnoEncoder",
     "SpectralConv1d",
+    "IntegralConvEmbedding",
 ]

reflectorch/models/encoders/integral_kernel_embedding.py

@@ -0,0 +1,390 @@
+from __future__ import annotations
+from typing import Union
+
+import torch
+from torch import nn, Tensor, stack, cat
+from reflectorch.models.activations import activation_by_name
+import reflectorch
+
+###embedding network adapted from the PANPE repository
+
+__all__ = [
+    "IntegralConvEmbedding",
+]
+
+class IntegralConvEmbedding(nn.Module):
+    def __init__(
+        self,
+        z_num: Union[int, tuple[int, ...]],
+        z_range: tuple[float, float] = None,
+        in_dim: int = 2,
+        kernel_coef: int = 16,
+        dim_embedding: int = 256,
+        conv_dims: tuple[int, ...] = (32, 64, 128),
+        num_blocks: int = 4,
+        use_batch_norm: bool = False,
+        use_layer_norm: bool = True,
+        use_fft: bool = False,
+        activation: str = "gelu",
+        conv_activation: str = "lrelu",
+        resnet_activation: str = "relu",
+    ) -> None:
+        super().__init__()
+
+        if isinstance(z_num, int):
+            z_num = (z_num,)
+        num_kernel = len(z_num)
+
+        if z_range is not None:
+            zs = [(z_range[0], z_range[1], nz) for nz in z_num]
+        else:
+            zs = z_num
+
+        self.in_dim = in_dim
+
+        self.kernels = nn.ModuleList(
+            [
+                IntegralKernelBlock(
+                    z,
+                    in_dim,
+                    kernel_coef=kernel_coef,
+                    latent_dim=dim_embedding,
+                    conv_dims=conv_dims,
+                    use_fft=use_fft,
+                    activation=activation,
+                    conv_activation=conv_activation,
+                )
+                for z in zs
+            ]
+        )
+
+        self.fc = reflectorch.models.networks.residual_net.ResidualMLP(
+            dim_in=dim_embedding * num_kernel, 
+            dim_out=dim_embedding,
+            layer_width=2 * dim_embedding,
+            num_blocks=num_blocks,
+            use_batch_norm=use_batch_norm,
+            use_layer_norm=use_layer_norm,
+            activation=resnet_activation,
+        )
+
+    def forward(self, q, y, drop_mask=None) -> Tensor:
+        x = cat([kernel(q, y, drop_mask=drop_mask) for kernel in self.kernels], dim=-1)
+        x = self.fc(x)
+        
+        return x
+
+
+class IntegralKernelBlock(nn.Module):
+    """
+    Examples:
+        >>> x = torch.rand(2, 100)
+        >>> y = torch.rand(2, 100, 3)
+        >>> block = IntegralKernelBlock((0, 1, 10), in_dim=3,  latent_dim=32)
+        >>> output = block(x, y)
+        >>> output.shape
+        torch.Size([2, 32])
+
+        >>> block = IntegralKernelBlock(10, in_dim=3,  latent_dim=32)
+        >>> output = block(x, y)
+        >>> output.shape
+        torch.Size([2, 32])
+    """
+
+    def __init__(
+        self,
+        z: tuple[float, float, int] or int,
+        in_dim: int,
+        kernel_coef: int = 2,
+        latent_dim: int = 32,
+        conv_dims: tuple[int, ...] = (32, 64, 128),
+        use_fft: bool = False,
+        activation: str = "gelu",
+        conv_activation: str = "lrelu",
+    ):
+        super().__init__()
+
+        if isinstance(z, int):
+            z_num = z
+            kernel = FullIntegralKernel(z_num, in_dim=in_dim, kernel_coef=kernel_coef)
+        else:
+            kernel = FastIntegralKernel(
+                z, in_dim=in_dim, kernel_coef=kernel_coef, activation=activation
+            )
+            z_num = z[-1]
+
+        assert z_num % 2 == 0, "z_num should be even"
+
+        self.kernel = kernel
+        self.z_num = z_num
+        self.in_dim = in_dim
+        self.latent_dim = latent_dim
+        self.use_fft = use_fft
+
+        self.fc_in_dim = self.latent_dim + self.in_dim * self.z_num
+        if self.use_fft:
+            self.fc_in_dim += self.in_dim * 2 + self.in_dim * self.z_num
+
+        self.conv = reflectorch.models.encoders.conv_encoder.ConvEncoder(
+            dim_avpool=8,
+            hidden_channels=conv_dims,
+            in_channels=in_dim,
+            dim_embedding=latent_dim,
+            activation=conv_activation,
+        )
+        self.fc = FCBlock(
+            in_dim=self.fc_in_dim, hid_dim=self.latent_dim * 2, out_dim=self.latent_dim
+        )
+
+    def forward(self, x: Tensor, y: Tensor, drop_mask: Tensor = None) -> Tensor:
+        x = self.kernel(x, y, drop_mask=drop_mask)
+
+        assert x.shape == (x.shape[0], self.in_dim, self.z_num)
+
+        xc = self.conv(x)  # (batch, latent_dim)
+
+        assert xc.shape == (x.shape[0], self.latent_dim)
+
+        if self.use_fft:
+            fft_x = torch.fft.rfft(x, dim=-1, norm="ortho")  # (batch, in_dim, z_num)
+
+            fft_x = torch.cat(
+                [fft_x.real, fft_x.imag], -1
+            )  # (batch, in_dim, 2 * z_num)
+
+            assert fft_x.shape == (x.shape[0], x.shape[1], self.z_num + 2)
+
+            fft_x = fft_x.flatten(1)  # (batch, in_dim * (z_num + 2))
+
+            x = torch.cat(
+                [x.flatten(1), fft_x, xc], -1
+            )  # (batch, in_dim * z_num * 3 + latent_dim)
+        else:
+            x = torch.cat([x.flatten(1), xc], -1)
+
+        assert (
+            x.shape[1] == self.fc_in_dim
+        ), f"Expected dim {self.fc_in_dim}, got {x.shape[1]}"
+
+        x = self.fc(x)  # (batch, latent_dim)
+
+        return x
+
+
+class FastIntegralKernel(nn.Module):
+    def __init__(
+        self,
+        z: tuple[float, float, int],
+        kernel_coef: int = 16,
+        in_dim: int = 1,
+        activation: str = "gelu",
+    ):
+        super().__init__()
+
+        z = torch.linspace(*z)
+
+        self.kernel = FCBlock(
+            in_dim + 2, kernel_coef * in_dim, in_dim, activation=activation
+        )
+
+        self.register_buffer("z", z)
+
+    def _get_z(self, x: Tensor):
+        # x.shape == (batch_size, num_x)
+        dz = self.z[1] - self.z[0]
+        indices = torch.ceil((x - self.z[0] - dz / 2) / dz).to(torch.int64)
+
+        z = torch.index_select(self.z, 0, indices.flatten()).view(*x.shape)
+
+        return z, indices
+
+    def forward(self, x: Tensor, y: Tensor, drop_mask=None):
+        z, indices = self._get_z(x)
+        xz = torch.stack([x, z], -1)
+        kernel_input = torch.cat([xz, y], -1)
+        output = self.kernel(kernel_input)  # (batch, x_num, in_dim)
+
+        output = compute_means(
+            output * y, indices, self.z.shape[-1], drop_mask=drop_mask
+        )  # (batch, z_num, in_dim)
+
+        output = output.swapaxes(1, 2)  # (batch, in_dim, z_num)
+
+        return output
+
+
+class FullIntegralKernel(nn.Module):
+    def __init__(
+        self,
+        z_num: int,
+        kernel_coef: int = 1,
+        in_dim: int = 1,
+    ):
+        super().__init__()
+
+        self.z_num = z_num
+        self.in_dim = in_dim
+
+        self.kernel = nn.Sequential(
+            nn.Linear(in_dim + 1, z_num * kernel_coef),
+            nn.LayerNorm(z_num * kernel_coef),
+            nn.ReLU(),
+            nn.Linear(z_num * kernel_coef, z_num * in_dim),
+        )
+
+    def forward(self, x: Tensor, y: Tensor, drop_mask=None):
+        # x.shape == (batch_size, num_x)
+        # y.shape == (batch_size, num_x, in_dim)
+        # drop_mask.shape == (batch_size, num_x)
+
+        batch_size, num_x = x.shape
+
+        kernel_input = torch.cat([x.unsqueeze(-1), y], -1)  # (batch, x_num, in_dim + 1)
+        x = self.kernel(kernel_input)  # (batch, x_num, z_num * in_dim)
+        x = x.reshape(
+            *x.shape[:-1], self.z_num, self.in_dim
+        )  # (batch, x_num, z_num, in_dim)
+        # permute to get (batch, z_num, x_num, in_dim)
+        x = x.permute(0, 2, 1, 3)
+
+        y = y.unsqueeze(1)  # (batch, 1, x_num, in_dim)
+
+        assert x.shape == (
+            batch_size,
+            self.z_num,
+            num_x,
+            self.in_dim,
+        )  # (batch, z_num, in_dim, x_num)
+        assert y.shape == (
+            batch_size,
+            1,
+            num_x,
+            self.in_dim,
+        )  # (batch, 1, x_num, in_dim)
+
+        if drop_mask is not None:
+            x = x * y
+            x = x.permute(0, 2, 1, 3)  # (batch, x_num, z_num, in_dim)
+            x = masked_mean(x, drop_mask)
+        else:
+            x = (x * y).mean(-2)  # (batch, z_num, in_dim)
+
+        assert x.shape == (batch_size, self.z_num, self.in_dim), f"{x.shape}"
+
+        x = x.swapaxes(1, 2)  # (batch, in_dim, z_num)
+
+        return x
+
+
+class FCBlock(nn.Module):
+    def __init__(
+        self,
+        in_dim: int = 2,
+        hid_dim: int = 16,
+        out_dim: int = 16,
+        activation: str = "gelu",
+    ):
+        super().__init__()
+
+        self.fc1 = nn.Linear(in_dim, hid_dim)
+        self.layer_norm = nn.LayerNorm(hid_dim)
+        self.activation = activation_by_name(activation)()
+        self.fc2 = nn.Linear(hid_dim, out_dim)
+
+    def forward(self, x: Tensor) -> Tensor:
+        x = self.fc1(x)
+        x = self.layer_norm(x)
+        x = self.activation(x)
+        x = self.fc2(x)
+        return x
+        # return self.kernel(x)
+
+
+def compute_means(x, indices, z: int, drop_mask: Tensor = None):
+    """
+    Compute the mean values of tensor 'x' for each unique index in 'indices' across each batch.
+
+    This function calculates the mean of elements in 'x' that correspond to each unique index in 'indices'.
+    The computation is performed for each batch separately, and the function is optimized to avoid Python loops
+    by using advanced PyTorch operations.
+
+    Parameters:
+    x (torch.Tensor): A tensor of shape (batch_size, n, d) containing the values to be averaged.
+                      'x' should be a floating-point tensor.
+    indices (torch.Tensor): An integer tensor of shape (batch_size, n) containing the indices.
+                            The values in 'indices' should be in the range [0, z-1].
+    z (int): The number of unique indices. This determines the second dimension of the output tensor.
+    drop_mask (torch.Tensor): A boolean tensor of shape (batch_size, n) containing a mask for the indices to drop.
+                              If None, all indices are used.
+
+    Returns:
+    torch.Tensor: A tensor of shape (batch_size, z, d) containing the mean values for each index in each batch.
+                  If an index does not appear in a batch, its corresponding mean values are zeros.
+
+    Example:
+    >>> batch_size, n, d, z = 3, 4, 5, 6
+    >>> indices = torch.randint(0, z, (batch_size, n))
+    >>> x = torch.randn(batch_size, n, d)
+    >>> y = compute_means(x, indices, z)
+    >>> print(y.shape)
+    torch.Size([3, 6, 5])
+    """
+
+    batch_size, n, d = x.shape
+    device = x.device
+
+    drop = drop_mask is not None
+
+    # Initialize tensors to hold sums and counts
+    sums = torch.zeros(batch_size, z + int(drop), d, device=device)
+    counts = torch.zeros(batch_size, z + int(drop), device=device)
+
+    if drop_mask is not None:
+        # Set the values of the indices to drop to z
+        indices = indices.masked_fill(~drop_mask, z)
+
+    indices_expanded = indices.unsqueeze(-1).expand_as(x)
+    sums.scatter_add_(1, indices_expanded, x)
+    counts.scatter_add_(1, indices, torch.ones_like(indices, dtype=x.dtype))
+
+    if drop:
+        # Remove the z values from the sums and counts
+        sums = sums[:, :-1]
+        counts = counts[:, :-1]
+
+    # Compute the mean and handle division by zero
+    mean = sums / counts.unsqueeze(-1).clamp(min=1)
+
+    return mean
+
+
+def masked_mean(x, mask):
+    """
+    Computes the mean of tensor x along the x_size dimension,
+    while masking out elements where the corresponding value in the mask is False.
+
+    Args:
+    x (torch.Tensor): A tensor of shape (batch, x_size, z, d).
+    mask (torch.Tensor): A boolean mask of shape (batch, x_size).
+
+    Returns:
+    torch.Tensor: The result tensor of shape (batch, z, d) after applying the mask and computing the mean.
+    """
+    if not mask.dtype == torch.bool:
+        raise TypeError("Mask must be a boolean tensor.")
+
+    # Ensure the mask is broadcastable to the shape of x
+    mask = mask.unsqueeze(-1).unsqueeze(-1)
+    masked_x = x * mask
+
+    # Compute the sum and the count of valid (unmasked) elements along the x_size dimension
+    sum_x = masked_x.sum(dim=1)
+    count_x = mask.sum(dim=1)
+
+    # Avoid division by zero
+    count_x[count_x == 0] = 1
+
+    # Compute the mean
+    mean_x = sum_x / count_x
+
+    return mean_x

reflectorch/models/networks/mlp_networks.py

@@ -7,6 +7,7 @@ from torch import nn, cat, split, Tensor
 
 from reflectorch.models.networks.residual_net import ResidualMLP
 from reflectorch.models.encoders.conv_encoder import ConvEncoder
+from reflectorch.models.encoders.integral_kernel_embedding import IntegralConvEmbedding
 from reflectorch.models.encoders.fno import FnoEncoder
 from reflectorch.models.activations import activation_by_name
 
@@ -18,7 +19,7 @@ class NetworkWithPriors(nn.Module):
         :align: center
 
     Args:
-        embedding_net_type (str): the type of embedding network, either 'conv' or 'fno'.
+        embedding_net_type (str): the type of embedding network, either 'conv', 'fno' or 'integral_conv'.
         embedding_net_kwargs (dict): dictionary containing the keyword arguments for the embedding network.
         dim_out (int, optional): the dimension of the output produced by the MLP. Defaults to 8.
         dim_conditioning_params (int, optional): the dimension of other parameters the network is conditioned on (e.g. for the smearing coefficient dq/q)
@@ -66,6 +67,8 @@ class NetworkWithPriors(nn.Module):
             self.embedding_net = ConvEncoder(**embedding_net_kwargs)
         elif embedding_net_type == 'fno':
             self.embedding_net = FnoEncoder(**embedding_net_kwargs)
+        elif embedding_net_type == 'integral_conv':
+            self.embedding_net = IntegralConvEmbedding(**embedding_net_kwargs)
         elif embedding_net_type == 'no_embedding_net':
             self.embedding_net = nn.Identity()
         else:
@@ -108,7 +111,7 @@ class NetworkWithPriors(nn.Module):
             self.embedding_net.load_weights(pretrained_embedding_net)
 
 
-    def forward(self, curves, bounds, q_values=None, conditioning_params=None):
+    def forward(self, curves, bounds, q_values=None, conditioning_params=None, key_padding_mask=None, unscaled_q_values=None):
         """
         Args:
             scaled_curves (torch.Tensor): Input tensor of shape [batch_size, n_points] or [batch_size, n_channels, n_points].
@@ -121,13 +124,16 @@ class NetworkWithPriors(nn.Module):
             curves = curves.unsqueeze(1)
 
         additional_channels = []
-        if q_values is not None:
+        if q_values is not None and not isinstance(self.embedding_net, IntegralConvEmbedding):
             additional_channels.append(q_values.unsqueeze(1))
 
         if additional_channels:
             curves = torch.cat([curves] + additional_channels, dim=1)  # [batch_size, n_channels, n_points]
 
-        x = self.embedding_net(curves)
+        if isinstance(self.embedding_net, IntegralConvEmbedding):
+            x = self.embedding_net(q=unscaled_q_values.float(), y=curves.permute(0, 2, 1), drop_mask=key_padding_mask)
+        else:
+            x = self.embedding_net(curves)
 
         if self.conditioning == 'concat': 
             x = torch.cat([x, bounds] + ([conditioning_params] if conditioning_params is not None else []), dim=-1)

reflectorch/runs/utils.py

@@ -278,7 +278,7 @@ def load_pretrained(model, model_name: str, saved_models_dir: Path):
         raise FileNotFoundError(f'File {str(model_path)} does not exist.')
     
     try:
-        pretrained = torch.load(model_path)
+        pretrained = torch.load(model_path, weights_only=False)
     except Exception as err:
         raise RuntimeError(f'Could not load model from {str(model_path)}') from err
     
@@ -295,6 +295,8 @@ def load_pretrained(model, model_name: str, saved_models_dir: Path):
 def init_dset(config: dict):
     """Initializes the dataset / dataloader object"""
     dset_cls = globals().get(config['cls']) if 'cls' in config else ReflectivityDataLoader
+    dset_kwargs = config.get('kwargs', {})
+
     prior_sampler = init_from_conf(config['prior_sampler'])
     intensity_noise = init_from_conf(config['intensity_noise'])
     q_generator = init_from_conf(config['q_generator'])
@@ -309,6 +311,7 @@ def init_dset(config: dict):
         curves_scaler=curves_scaler,
         smearing=smearing,
         q_noise=q_noise,
+        **dset_kwargs,
     )
 
     return dset
@@ -358,7 +361,7 @@ def convert_pt_to_safetensors(input_dir):
                 continue
 
             print(f"Converting {pt_file_path} to .safetensors format.")
-            data_pt = torch.load(pt_file_path)
+            data_pt = torch.load(pt_file_path, weights_only=False)
             model_state_dict = data_pt["model"]
             model_state_dict = split_complex_tensors(model_state_dict) #handle tensors with complex dtype which are not natively supported by safetensors
 

reflectorch/utils.py

@@ -66,3 +66,33 @@ def energy_to_wavelength(energy: float):
 def wavelength_to_energy(wavelength: float):
     """Conversion from photon wavelength (angstroms) to photon energy (eV)"""
     return 1.2398 / wavelength * 1e4
+
+def get_filtering_mask(Q, R, dR, threshold=0.3, consecutive=3,
+                        remove_singles=True, remove_consecutives=True,
+                        q_start_trunc=0.1):
+    Q, R, dR = Q.copy(), R.copy(), dR.copy()
+    rel_error = np.abs(dR / R)
+
+    # Mask for singles
+    mask_singles = (rel_error >= threshold) if remove_singles else np.zeros_like(Q, dtype=bool)
+
+    # Mask for truncation
+    mask_consecutive = np.zeros_like(Q, dtype=bool)
+    if remove_consecutives:
+        count = 0
+        cutoff_idx = None
+        for i in range(len(Q)):
+            if Q[i] < q_start_trunc:
+                continue
+            if rel_error[i] >= threshold:
+                count += 1
+                if count >= consecutive:
+                    cutoff_idx = i - consecutive + 1
+                    break
+            else:
+                count = 0
+        if cutoff_idx is not None:
+            mask_consecutive[cutoff_idx:] = True
+
+    final_mask = mask_singles | mask_consecutive
+    return ~final_mask

{reflectorch-1.3.0.dist-info → reflectorch-1.4.0.dist-info}/METADATA

@@ -1,10 +1,9 @@
 Metadata-Version: 2.4
 Name: reflectorch
-Version: 1.3.0
+Version: 1.4.0
 Summary: A Pytorch-based package for the analysis of reflectometry data
 Author-email: Vladimir Starostin <vladimir.starostin@uni-tuebingen.de>, Valentin Munteanu <valentin.munteanu@uni-tuebingen.de>
 Maintainer-email: Valentin Munteanu <valentin.munteanu@uni-tuebingen.de>, Vladimir Starostin <vladimir.starostin@uni-tuebingen.de>, Alexander Hinderhofer <alexander.hinderhofer@uni-tuebingen.de>
-License-Expression: MIT
 Project-URL: Source, https://github.com/schreiber-lab/reflectorch/
 Project-URL: Issues, https://github.com/schreiber-lab/reflectorch/issues
 Project-URL: Documentation, https://schreiber-lab.github.io/reflectorch/
@@ -106,6 +105,8 @@ Configuration files and the corresponding pretrained model weights are hosted on
 <!-- [![Docker](https://img.shields.io/badge/Docker-2496ED.svg?style=flat&logo=docker&logoColor=white)](https://hub.docker.com/)
 Docker images for reflectorch *will* be hosted on Dockerhub. -->
 
+## Contributing
+If you'd like to contribute to the package, please see our [Contributing Guidelines](CONTRIBUTING.md) for details.
 
 ## Citation
 If you find our work useful in your research, please cite as follows: