def _cartesian_mask()

in activemri/baselines/simple_baselines.py [0:0]

• 23 lines of code
• 6 McCabe index (conditional complexity)

    def _cartesian_mask(self, current_mask: np.ndarray) -> np.ndarray:
        batch_size, image_width = current_mask.shape
        pdf_x = RandomLowBiasPolicy._normal_pdf(
            image_width, 0.5 / (image_width / 10.0) ** 2
        )
        pdf_x = np.expand_dims(pdf_x, axis=0)
        lmda = image_width / (2.0 * self.acceleration)
        # add uniform distribution
        pdf_x += lmda * 1.0 / image_width
        # remove previously chosen columns
        # note that pdf_x designed for centered masks
        new_mask = (
            np.fft.ifftshift(current_mask, axes=1)
            if not self.centered
            else current_mask.copy()
        )
        pdf_x = pdf_x * np.logical_not(new_mask)
        # normalize probabilities and choose accordingly
        pdf_x /= pdf_x.sum(axis=1, keepdims=True)
        indices = [
            self.rng.choice(image_width, 1, False, pdf_x[i]).item()
            for i in range(batch_size)
        ]
        new_mask[range(batch_size), indices] = 1
        if not self.centered:
            new_mask = np.fft.ifftshift(new_mask, axes=1)
        return new_mask