text:
import math
from enum import Enum
from typing import List, Tuple, Optional, Dict

import torch
from torch import Tensor

if __name__ == "__main__":
    from torchvision_0_11_1_transforms import functional as F
    from torchvision_0_11_1_transforms.functional import InterpolationMode
else:
    from .torchvision_0_11_1_transforms import functional as F
    from .torchvision_0_11_1_transforms.functional import InterpolationMode

#__all__ = ["AutoAugmentPolicy", "AutoAugment", "RandAugment", "TrivialAugmentWide"]
__all__ = ["RandAugment"]

def _apply_op(
    img: Tensor, op_name: str, magnitude: float, interpolation: InterpolationMode, fill: Optional[List[float]]
):
    if op_name == "ShearX":
        img = F.affine(
            img,
            angle=0.0,
            translate=[0, 0],
            scale=1.0,
            shear=[math.degrees(magnitude), 0.0],
            interpolation=interpolation,
            fill=fill,
        )
    elif op_name == "ShearY":
        img = F.affine(
            img,
            angle=0.0,
            translate=[0, 0],
            scale=1.0,
            shear=[0.0, math.degrees(magnitude)],
            interpolation=interpolation,
            fill=fill,
        )
    elif op_name == "TranslateX":
        img = F.affine(
            img,
            angle=0.0,
            translate=[int(magnitude), 0],
            scale=1.0,
            interpolation=interpolation,
            shear=[0.0, 0.0],
            fill=fill,
        )
    elif op_name == "TranslateY":
        img = F.affine(
            img,
            angle=0.0,
            translate=[0, int(magnitude)],
            scale=1.0,
            interpolation=interpolation,
            shear=[0.0, 0.0],
            fill=fill,
        )
    elif op_name == "Rotate":
        img = F.rotate(img, magnitude, interpolation=interpolation, fill=fill)
    elif op_name == "Brightness":
        img = F.adjust_brightness(img, 1.0 + magnitude)
    elif op_name == "Color":
        img = F.adjust_saturation(img, 1.0 + magnitude)
    elif op_name == "Contrast":
        img = F.adjust_contrast(img, 1.0 + magnitude)
    elif op_name == "Sharpness":
        img = F.adjust_sharpness(img, 1.0 + magnitude)
    elif op_name == "Posterize":
        img = F.posterize(img, int(magnitude))
    elif op_name == "Solarize":
        img = F.solarize(img, magnitude)
    elif op_name == "AutoContrast":
        img = F.autocontrast(img)
    elif op_name == "Equalize":
        img = F.equalize(img)
    elif op_name == "Invert":
        img = F.invert(img)
    elif op_name == "Identity":
        pass
    else:
        raise ValueError(f"The provided operator {op_name} is not recognized.")
    return img

class RandAugment(torch.nn.Module):
    r"""RandAugment data augmentation method based on
    `"RandAugment: Practical automated data augmentation with a reduced search space"
    <https://arxiv.org/abs/1909.13719>`_.
    If the image is torch Tensor, it should be of type torch.uint8, and it is expected
    to have [..., 1 or 3, H, W] shape, where ... means an arbitrary number of leading dimensions.
    If img is PIL Image, it is expected to be in mode "L" or "RGB".
    Args:
        num_ops (int): Number of augmentation transformations to apply sequentially.
        magnitude (int): Magnitude for all the transformations.
        num_magnitude_bins (int): The number of different magnitude values.
        interpolation (InterpolationMode): Desired interpolation enum defined by
            :class:`torchvision.transforms.InterpolationMode`. Default is ``InterpolationMode.NEAREST``.
            If input is Tensor, only ``InterpolationMode.NEAREST``, ``InterpolationMode.BILINEAR`` are supported.
        fill (sequence or number, optional): Pixel fill value for the area outside the transformed
            image. If given a number, the value is used for all bands respectively.
    """

    def __init__(
        self,
        num_ops: int = 2,
        magnitude: int = 9,
        num_magnitude_bins: int = 31,
        interpolation: InterpolationMode = InterpolationMode.NEAREST,
        fill: Optional[List[float]] = None,
    ) -> None:
        super().__init__()
        self.num_ops = num_ops
        self.magnitude = magnitude
        self.num_magnitude_bins = num_magnitude_bins
        self.interpolation = interpolation
        self.fill = fill

    def _augmentation_space(self, num_bins: int, image_size: List[int]) -> Dict[str, Tuple[Tensor, bool]]:
        return {
            # op_name: (magnitudes, signed)
            "Identity": (torch.tensor(0.0), False),
            "ShearX": (torch.linspace(0.0, 0.3, num_bins), True),
            "ShearY": (torch.linspace(0.0, 0.3, num_bins), True),
            "TranslateX": (torch.linspace(0.0, 150.0 / 331.0 * image_size[0], num_bins), True),
            "TranslateY": (torch.linspace(0.0, 150.0 / 331.0 * image_size[1], num_bins), True),
            "Rotate": (torch.linspace(0.0, 30.0, num_bins), True),
            "Brightness": (torch.linspace(0.0, 0.9, num_bins), True),
            "Color": (torch.linspace(0.0, 0.9, num_bins), True),
            "Contrast": (torch.linspace(0.0, 0.9, num_bins), True),
            "Sharpness": (torch.linspace(0.0, 0.9, num_bins), True),
            "Posterize": (8 - (torch.arange(num_bins) / ((num_bins - 1) / 4)).round().int(), False),
            "Solarize": (torch.linspace(255.0, 0.0, num_bins), False),
            "AutoContrast": (torch.tensor(0.0), False),
            "Equalize": (torch.tensor(0.0), False),
        }

    def forward(self, img: Tensor, label: Tensor) -> tuple:
        """
            img (PIL Image or Tensor): Image to be transformed.
            label (PIL Image or Tensor): Label to be transformed.
        Returns:
            PIL Image or Tensor: Transformed image & label.
        """
        fill = self.fill
        if isinstance(img, Tensor):
            if isinstance(fill, (int, float)):
                fill = [float(fill)] * F.get_image_num_channels(img)
            elif fill is not None:
                fill = [float(f) for f in fill]

        for _ in range(self.num_ops):
            op_meta = self._augmentation_space(self.num_magnitude_bins, F.get_image_size(img))
            op_index = int(torch.randint(len(op_meta), (1,)).item())
            op_name = list(op_meta.keys())[op_index]
            magnitudes, signed = op_meta[op_name]
            magnitude = float(magnitudes[self.magnitude].item()) if magnitudes.ndim > 0 else 0.0
            if signed and torch.randint(2, (1,)):
                magnitude *= -1.0
            img = _apply_op(img, op_name, magnitude, interpolation=self.interpolation, fill=fill)
            if op_name in ["Identity", "ShearX", "ShearY", "TranslateX", "TranslateY", "Rotate"]:
                label = _apply_op(label, op_name, magnitude, interpolation=self.interpolation, fill=fill)
        return img, label

    def __repr__(self) -> str:
        s = self.__class__.__name__ + "("
        s += "num_ops={num_ops}"
        s += ", magnitude={magnitude}"
        s += ", num_magnitude_bins={num_magnitude_bins}"
        s += ", interpolation={interpolation}"
        s += ", fill={fill}"
        s += ")"
        return s.format(**self.__dict__)

if __name__ == "__main__":
    from PIL import Image
    import matplotlib.pyplot as plt
    import numpy as np
    root = "G:/Codes/RealTime-Segementation/datasets/VOC2012"
    img_path = f"{root}/JPEGImages/2007_000033.jpg"
    label_path = f"{root}/SegmentationObject/2007_000033.png"
    img = Image.open(img_path)
    label = Image.open(label_path).convert('RGB')
    plt.subplot(221), plt.title("Ori Image"), plt.imshow(np.asarray(img))
    plt.subplot(222), plt.title("Ori Label"), plt.imshow(np.asarray(label))
    transforms = RandAugment(magnitude=15)
    img, label = transforms(img, label)
    plt.subplot(223), plt.title("Ori Image"), plt.imshow(np.asarray(img))
    plt.subplot(224), plt.title("Ori Label"), plt.imshow(np.asarray(label))
    plt.show()


text:
import sys
from utils.train_options import TrainOptions
from networks.kd_model import NetModel
from networks.Seg_kd_model import Seg_NetModel
import logging
import warnings
warnings.filterwarnings("ignore")
from torch.utils import data
from dataset.datasets import CSDataSet
import numpy as np
from tensorboardX import SummaryWriter

LOG_DIR = './logs/struct_dis_res50_to_res18_0.5_256_256'
IMG_MEAN = np.array((104.00698793,116.66876762,122.67891434), dtype=np.float32)
args = TrainOptions().initialize()
h, w = map(int, args.input_size.split(','))

trainloader = data.DataLoader(CSDataSet(args.data_dir, '/home/remo/Desktop/structure_knowledge_distillation-master/dataset/demo/tran.txt', max_iters=args.num_steps*args.batch_size, crop_size=(h, w),
                scale=args.random_scale, mirror=args.random_mirror, mean=IMG_MEAN),
                batch_size=args.batch_size, shuffle=True, num_workers=4, pin_memory=True)
valloader = data.DataLoader(CSDataSet(args.data_dir, '/home/remo/Desktop/structure_knowledge_distillation-master/dataset/demo/tran.txt', crop_size=(256, 256), mean=IMG_MEAN, scale=False, mirror=False),
                                batch_size=1, shuffle=False, pin_memory=True)
# save_steps = int(2975/args.batch_size)
save_steps = 1
model = Seg_NetModel(args)
tblogger = SummaryWriter(LOG_DIR)
itr = 0
mean_IU, IU_array = 0,0
for epoch in range(args.start_epoch, args.epoch_nums):
    for step, data in enumerate(trainloader, args.last_step+1):
        '''
        data:[image,mask,img_shape,img_name]
        image:[B, N, W, H] torch.tensor
        mask:[B, W, H] torch.tensor
        img_shape:[[H, W, C],
                    [H, W, C]
                    ...
                    ] Batch_size个 [B,3]
        img_name:['img_name1',
                    'img_name2'] [B,1]
        '''
        model.adjust_learning_rate(args.lr_g, model.G_solver, step)
        model.adjust_learning_rate(args.lr_d, model.D_solver, step)
        model.set_input(data)
        model.optimize_parameters()
        model.print_info(epoch, step)
        # if (step > 1) and ((step % save_steps == 0) and (step > args.num_steps - 1000)) or (step == args.num_steps - 1):
        if (step > 1) and (step % save_steps == 0):
            mean_IU, IU_array = model.evalute_model(model.student, valloader, '0', '256,256', 9, True)
            model.save_ckpt(epoch, step, mean_IU, IU_array)
            logging.info('[val 512,512] mean_IU:{:.6f}  IU_array:{}'.format(mean_IU, IU_array))
            tblogger.add_scalar('mIoU', mean_IU, itr)
        itr += 1

ext: py

reason: import numpy


text:
"""Containers and aggregation functions for sufficient statistics."""

import numpy as np
from scipy import sparse
from scipy.sparse import csr_matrix, hstack

class EntityStatistics:
    """Entity-level sufficient statistics."""
    def __init__(self, counts, n_observations):
        self.counts = counts
        self.n_observations = n_observations

    @classmethod
    def from_observations(cls, observations, labels, n_classes=None):
        """Aggregate observations to entity level.

        Parameters
        ----------
        observations : scipy.sparse.csr_matrix
            binary document-term matrix
        labels : numpy.ndarray
            class labels. Must have y.size == X.shape[0]
        n_classes : int, optional
            number of classes (defaults to `y.max() + 1`)

        Returns
        -------
        EntityStatistics
        """
        if labels.size != observations.shape[0]:
            raise ValueError('Incompatible shapes.')
        n = labels.size
        if n_classes is None:
            n_classes = labels.max() + 1
        dtype = np.int32
        data, row, col = np.ones(n, dtype=dtype), labels.astype(dtype), np.arange(n, dtype=dtype)
        assignments_mat = csr_matrix((data, (row, col)), shape=(n_classes, observations.shape[0]))
        counts = assignments_mat.dot(observations)
        n_observations = np.array(assignments_mat.sum(axis=1)).ravel()

        return cls(counts, n_observations)

    @property
    def n_tokens(self):
        """Get number of distinct tokens."""
        return self.counts.shape[1]

    @property
    def size(self):
        """Get number of entities."""
        return self.n_observations.size

    def __repr__(self):
        s = 'EntityStatistics({} entities, {} observations)'
        return s.format(self.counts.shape[0], self.n_observations.sum())

    def extend_tokens(self, n_tokens):
        """Increase the number of tokens.

        New tokens are appended at the end with all counts
        set to zero.

        Parameters
        ----------
        n_tokens : int >= 0
            Number of tokens to add.
        """
        if n_tokens == 0:
            return self
        new_block = csr_matrix((self.counts.shape[0], n_tokens), dtype=self.counts.dtype)
        self.counts = hstack([self.counts, new_block])
        return self

    def add(self, other):
        """add sufficient statistics."""
        if not other.counts.shape == self.counts.shape:
            raise ValueError('Incompatible shapes.')
        self.counts += other.counts
        self.n_observations += other.n_observations
        return self

    def copy(self):
        """return new EntityStatistics with copied parameter arrays."""
        return self.__class__(self.counts.copy(), self.n_observations.copy())

class TokenStatistics:
    """Container for count aggregation to token level.

    The constructor is not meant to be used directly.
    New instances are created using the class methods

    `TokenStatistics.from_entity_statistics`
    `TokenStatistics.from_observations`
    """
    def __init__(self, total_weights, positive_weights, negative_weights):
        self.total_weights = total_weights
        self.positive_weights = sparse.csr_matrix(positive_weights)
        self.negative_weights = sparse.csr_matrix(negative_weights)

    @classmethod
    def from_entity_statistics(cls, entity_stats):
        """Aggregate EntityStatistics to token level."""
        _, n_tokens = entity_stats.counts.shape
        n_obs_vals, n_obs_cnts = np.unique(entity_stats.n_observations, return_counts=True)
        max_n_obs_count = n_obs_vals[-1]
        counts = entity_stats.counts.tocoo()

        # aggregate positive counts directly
        # cases with zero positive count are skipped
        # but can be recovered by summing over rows
        row = counts.col
        col = counts.data
        data = np.ones_like(col)
        positive_weights = sparse.coo_matrix(
            (data, (row, col)),
            shape=(n_tokens, max_n_obs_count + 1)
        )

        # aggregate negative counts indirectly
        # 1. Create default case from n_obs_vals, n_obs_counts
        row = [np.repeat(np.arange(n_tokens), n_obs_vals.size)]
        col = [np.tile(n_obs_vals, n_tokens)]
        data = [np.tile(n_obs_cnts, n_tokens)]

        # 2. Subtract cases that have any observations
        row.append(counts.col)
        col.append(entity_stats.n_observations[counts.row])
        data.append(np.full_like(counts.row, -1))

        # 3. Add cases with observations with correct counts
        row.append(counts.col)
        col.append(entity_stats.n_observations[counts.row] - counts.data)
        data.append(np.ones_like(counts.row))

        negative_weights = sparse.coo_matrix(
            (np.concatenate(data), (np.concatenate(row), np.concatenate(col))),
            shape=(n_tokens, max_n_obs_count + 1)
        )

        total_weights = np.zeros(n_obs_vals[-1] + 1, counts.dtype)
        total_weights[n_obs_vals] = n_obs_cnts
        # backward compatibility: convert to dense format
        return cls(total_weights, positive_weights, negative_weights)

    @classmethod
    def from_observations(cls, observations):
        """Aggregate observations to token level.

        Parameters
        ----------
        observations: scipy.sparse.csr_matrix
            Binary matrix with token counts per observation.
        """
        dtype = dict(dtype=np.int32)
        n_rows, n_tokens = observations.shape
        total_weights = np.array([0, n_rows], **dtype)

        counts = np.array(observations.sum(axis=0)).ravel()
        positive_weights = sparse.coo_matrix(
            (counts, (np.arange(n_tokens, **dtype), np.ones(n_tokens, **dtype))),
            shape=(n_tokens, 2)
        )
        negative_weights = sparse.coo_matrix(
            (n_rows - counts, (np.arange(n_tokens, **dtype), np.ones(n_tokens, **dtype))),
            shape=(n_tokens, 2)
        )

        return cls(total_weights, positive_weights, negative_weights)

    def add(self, other):
        """Add token counts.

        Parameters
        ----------
        other: ptfidf.aggregation.TokenStatistics
            Statistics to add.

        Returns
        -------
        ptfidf.aggregation.TokenStatistics
            New object with added statistics.
        """
        dtype = dict(dtype=np.int32)
        if self.size != other.size:
            raise ValueError('Incompatible number of tokens: %d != %d' % (self.size, other.size))
        n_tokens = self.size
        n_max = 1 + max(self.max_count, other.max_count)
        total_weights = np.zeros(n_max, **dtype)
        positive_weights = sparse.coo_matrix((n_tokens, n_max), **dtype)
        negative_weights = sparse.coo_matrix((n_tokens, n_max), **dtype)
        for obj in [self, other]:
            total_weights[:obj.total_weights.size] += obj.total_weights
            padding = sparse.coo_matrix(
                (n_tokens, n_max - obj.max_count - 1),
                dtype=positive_weights.dtype
            )
            positive_weights += sparse.hstack([obj.positive_weights, padding])
            negative_weights += sparse.hstack([obj.negative_weights, padding])
        return self.__class__(total_weights, positive_weights, negative_weights)

    def extend_tokens(self, n_tokens: int) -> None:
        """Increase the number of tokens.

        New tokens are appended at the end with all counts
        set to zero.

        Parameters
        ----------
        n_tokens : int >= 0
            Number of tokens to add.

        Returns
        -------
        ptfidf.aggregation.TokenStatistics
            New instance with extended token number. All weight
            arrays are padded with zeros to the right.
        """
        if n_tokens == 0:
            return self.copy()
        positive_padding = sparse.coo_matrix(
            (n_tokens, self.positive_weights.shape[1]),
            dtype=self.positive_weights.dtype
        )
        negative_padding = sparse.coo_matrix(
            np.repeat(self.total_weights[None, :], n_tokens, axis=0)
        )
        return self.__class__(
            self.total_weights.copy(),
            sparse.vstack([self.positive_weights, positive_padding]),
            sparse.vstack([self.negative_weights, negative_padding])
        )

    def copy(self):
        """Create new instance with copied data."""
        return self.__class__(
            self.total_weights.copy(),
            self.positive_weights.copy(),
            self.negative_weights.copy()
        )

    def get_unique_weights(self):
        """Return deduplicated weights.

        Returns
        -------
        (positive_weights, negative_weights, index, inverse, counts) : tuple
            positive_weights and negative_weights are deduplicated.
            index, inverse, counts are defined like in numpy.unique
        """
        n = self.total_weights.size
        # Need to find simultaneous unique rows of positive_weights and negative weights
        projection = np.random.uniform(-1., 1., size=(n, min(n, 10)))
        hashed = np.c_[self.positive_weights.dot(projection), self.negative_weights.dot(projection)]
        _, index, inverse, counts = np.unique(
            hashed, axis=0,
            return_counts=True, return_index=True, return_inverse=True
        )
        return self.positive_weights[index], self.negative_weights[index], index, inverse, counts

    def __repr__(self):
        return 'TokenStatistics(%d tokens)' % self.size

    @property
    def size(self) -> int:
        """Get number of tokens."""
        return self.positive_weights.shape[0]

    @property
    def max_count(self) -> int:
        """Get max number of entity observations."""
        return self.total_weights.size - 1

max_stars_repo_name: mossadnik/probabilistic-tf-idf

max_issues_repo_name: mossadnik/probabilistic-tf-idf

max_forks_repo_name: mossadnik/probabilistic-tf-idf

reason: import numpy,from scipy


text:
import sys
from pathlib import Path

import matplotlib as mpl
import matplotlib.cm as cm
import numpy as np
import PIL.Image as pil
import torch
from torchvision import transforms

import networks

def run():
    input_name = sys.argv[1]
    output_name = sys.argv[2]
    colormap_name = sys.argv[3]

    if torch.cuda.is_available():
        device = torch.device("cuda")
    else:
        device = torch.device("cpu")

    root = Path(sys.path[0]).resolve()

    model_path = root / "models"
    encoder_path = str(model_path / "encoder.pth")
    depth_decoder_path = str(model_path / "depth.pth")

    # LOADING PRETRAINED MODEL
    encoder = networks.ResnetEncoder(18, False)
    loaded_dict_enc = torch.load(encoder_path, map_location=device)

    # extract the height and width of image that this model was trained with
    feed_height = loaded_dict_enc['height']
    feed_width = loaded_dict_enc['width']
    filtered_dict_enc = {
        k: v for k, v in loaded_dict_enc.items() if k in encoder.state_dict()}
    encoder.load_state_dict(filtered_dict_enc)
    encoder.to(device)
    encoder.eval()

    depth_decoder = networks.DepthDecoder(
        num_ch_enc=encoder.num_ch_enc, scales=range(4))

    loaded_dict = torch.load(depth_decoder_path, map_location=device)
    depth_decoder.load_state_dict(loaded_dict)

    depth_decoder.to(device)
    depth_decoder.eval()

    data_path = root.parent.parent / "data"
    input_path = str(data_path / "inputs" / input_name)
    output_path = str(data_path / "outputs" / output_name)

    with torch.no_grad():
        # Load image and preprocess
        input_image = pil.open(input_path).convert('RGB')
        original_width, original_height = input_image.size
        input_image = input_image.resize(
            (feed_width, feed_height), pil.LANCZOS)
        input_image = transforms.ToTensor()(input_image).unsqueeze(0)

        # PREDICTION
        input_image = input_image.to(device)
        features = encoder(input_image)
        outputs = depth_decoder(features)

        disp = outputs[("disp", 0)]
        disp_resized = torch.nn.functional.interpolate(
            disp, (original_height, original_width), mode="bilinear", align_corners=False)

        # Saving colormapped depth image
        prediction = disp_resized.squeeze().cpu().numpy()
        vmax = np.percentile(prediction, 95)
        normalizer = mpl.colors.Normalize(
            vmin=prediction.min(), vmax=vmax)
        mapper = cm.ScalarMappable(norm=normalizer, cmap=colormap_name)
        colormapped_im = (mapper.to_rgba(prediction)[
                          :, :, :3] * 255).astype(np.uint8)
        im = pil.fromarray(colormapped_im)
        im.save(output_path)

    print('Done')

if __name__ == '__main__':
    run()

max_stars_repo_name: ilharp/depth-estimate-gui

max_issues_repo_name: ilharp/depth-estimate-gui

max_forks_repo_name: depth-estimate-gui/depth-estimate-gui


text:
from __future__ import annotations

from collections.abc import Iterable
from collections.abc import Sequence
from functools import reduce

import numpy as np
from cachetools import cached
from cachetools.keys import hashkey
from scipy.linalg import eig
from scipy.linalg import eigvals
from scipy.linalg import expm
from scipy.linalg import inv

from chemex.nmr.constants import Distribution
from chemex.nmr.liouvillian import LiouvillianIS

def calculate_spectrometers(
    liouv: np.ndarray, delays: float | Iterable[float], dephasing: bool = False
) -> np.ndarray:
    delays_ = np.asarray(delays).reshape(-1)
    shape = liouv.shape
    spectrometer_list = []
    for a_liouvillian in liouv.reshape(-1, *shape[-2:]):
        s, vr = eig(a_liouvillian)  # type: ignore
        vri = inv(vr)
        if dephasing:
            indexes = np.where(abs(s.imag) < 1e-6)[0]  # type: ignore
            vr, s, vri = vr[:, indexes], s[indexes], vri[indexes, :]
        d = np.asarray([np.diag(np.exp(s * t)) for t in delays_])
        spectrometer_list.append((vr @ d @ vri).real)
    spectrometers = np.asarray(spectrometer_list).swapaxes(0, 1).reshape(-1, *shape)
    if spectrometers.shape[0] == 1:
        spectrometers = spectrometers[0]
    return spectrometers

def _get_key(liouvillian: LiouvillianIS, *args, **kwargs):
    return hashkey(liouvillian.basis, *args, **kwargs)

@cached(cache=dict(), key=_get_key)
def _make_perfect180(liouvillian: LiouvillianIS, spin: str) -> np.ndarray:
    size = liouvillian.size
    identity = np.eye(size).reshape((1, 1, size, size))
    compx, compy, compz = (f"{spin}{axis}" for axis in "xyz")
    perfect180 = {comp: identity.copy() for comp in (compx, compy)}
    for comp in liouvillian.basis.components:
        vect = liouvillian.vectors[comp].ravel()
        if compx in comp or compz in comp:
            perfect180[compy] -= 2 * np.diag(vect)
        if compy in comp or compz in comp:
            perfect180[compx] -= 2 * np.diag(vect)
    p180 = [perfect180[comp] for comp in (compx, compy, compx, compy)]
    return np.array(p180)

@cached(cache=dict(), key=_get_key)
def _make_perfect90(liouvillian: LiouvillianIS, spin: str) -> np.ndarray:
    size = liouvillian.size
    zeros = np.zeros((size, size))
    rot = liouvillian.matrices.get(f"b1x_{spin}", zeros)
    return np.asarray(expm(0.25 * rot)).reshape((1, 1, size, size))

@cached(cache=dict(), key=_get_key)
def _get_phases(liouvillian: LiouvillianIS) -> dict[str, np.ndarray]:
    phases = {}
    size = liouvillian.size
    zeros = np.zeros((size, size))
    for spin in "is":
        l_rotz = liouvillian.matrices.get(f"rotz_{spin}", zeros)
        phases[spin] = np.array([expm(n * 0.5 * np.pi * l_rotz) for n in range(4)])
    return phases

class Spectrometer:
    def keep(self, magnetization: np.ndarray, components: Iterable[str]) -> np.ndarray:
        return self.liouvillian.keep(magnetization, components)

    def _add_phases(self, propagator: np.ndarray, spin: str = "i") -> np.ndarray:
        phases = self._phases[spin]
        return np.array([phases[i] @ propagator @ phases[-i] for i in range(4)])

    def __init__(self, liouvillian: LiouvillianIS) -> None:
        self.liouvillian = liouvillian
        size = self.liouvillian.size
        self.identity = np.eye(self.liouvillian.size).reshape((1, 1, size, size))
        self._phases = _get_phases(liouvillian)
        self.perfect90_i = self._add_phases(_make_perfect90(liouvillian, "i"))
        self.perfect180_i = _make_perfect180(liouvillian, "i")
        self.perfect180_s = _make_perfect180(liouvillian, "s")
        self.zfilter = self.keep(
            self.identity,
            components={"ie", "se", "iz", "sz", "2izsz"}
            & set(liouvillian.basis.components),
        )
        self.calculate_i_flag = True
        self.calculate_s_flag = True
        self._pw90_i = 0.0
        self._pw90_s = 0.0
        self._p90_i = np.array(0.0)
        self._p180_i = np.array(0.0)
        self._p240_i = np.array(0.0)
        self._p180_s = np.array(0.0)

    def update(self, par_values: dict[str, float]) -> None:
        self.liouvillian.update(par_values)
        self.calculate_i_flag = True
        self.calculate_s_flag = True

    @property
    def par_values(self) -> dict[str, float]:
        return self.liouvillian.par_values

    @property
    def carrier_i(self) -> float:
        return self.liouvillian.carrier_i

    @carrier_i.setter
    def carrier_i(self, value: float):
        self.liouvillian.carrier_i = value
        self.calculate_i_flag = True
        self.calculate_s_flag = True

    @property
    def carrier_s(self) -> float:
        return self.liouvillian.carrier_s

    @carrier_s.setter
    def carrier_s(self, value: float):
        self.liouvillian.carrier_s = value
        self.calculate_i_flag = True
        self.calculate_s_flag = True

    @property
    def offset_i(self) -> float:
        return self.liouvillian.offset_i

    @offset_i.setter
    def offset_i(self, value: float):
        self.liouvillian.offset_i = value
        self.calculate_i_flag = True
        self.calculate_s_flag = True

    @property
    def offset_s(self) -> float:
        return self.liouvillian.offset_s

    @offset_s.setter
    def offset_s(self, value: float):
        self.liouvillian.offset_s = value
        self.calculate_i_flag = True
        self.calculate_s_flag = True

    @property
    def b1_i(self) -> float:
        return self.liouvillian.b1_i

    @b1_i.setter
    def b1_i(self, value: float):
        self._pw90_i = 1.0 / (4.0 * value) if value else 0.0
        self.liouvillian.b1_i = value
        self.calculate_i_flag = True
        self.calculate_s_flag = True

    @property
    def b1_s(self) -> float:
        return self.liouvillian.b1_s

    @b1_s.setter
    def b1_s(self, value: float):
        self._pw90_s = 1.0 / (4.0 * value) if value else 0.0
        self.liouvillian.b1_s = value
        self.calculate_i_flag = True
        self.calculate_s_flag = True

    @property
    def b1_i_inh_scale(self) -> float:
        return self.liouvillian.b1_i_inh_scale

    @b1_i_inh_scale.setter
    def b1_i_inh_scale(self, value: float):
        self.liouvillian.b1_i_inh_scale = value
        self.calculate_i_flag = True
        self.calculate_s_flag = True

    @property
    def b1_i_inh_res(self) -> int:
        return self.liouvillian.b1_i_inh_res

    @b1_i_inh_res.setter
    def b1_i_inh_res(self, value: int):
        self.liouvillian.b1_i_inh_res = value
        self.calculate_i_flag = True
        self.calculate_s_flag = True

    @property
    def jeff_i(self) -> Distribution:
        return self.liouvillian.jeff_i

    @jeff_i.setter
    def jeff_i(self, value: Distribution):
        self.liouvillian.jeff_i = value
        self.calculate_i_flag = True
        self.calculate_s_flag = True

    def get_equilibrium(self) -> np.ndarray:
        return self.liouvillian.get_equilibrium()

    def get_start_magnetization(
        self, terms: Iterable[str], atom: str = "h"
    ) -> np.ndarray:
        return self.liouvillian.get_start_magnetization(terms=terms, atom=atom)

    @property
    def detection(self) -> str:
        return self.liouvillian.detection

    @detection.setter
    def detection(self, value: str):
        self.liouvillian.detection = value

    def detect(self, magnetization: np.ndarray) -> float:
        return self.liouvillian.detect(magnetization)

    @property
    def gradient_dephasing(self) -> float:
        return self.liouvillian.gradient_dephasing

    @gradient_dephasing.setter
    def gradient_dephasing(self, value: float):
        self.liouvillian.gradient_dephasing = value
        self.calculate_i_flag = True
        self.calculate_s_flag = True

    def delays(self, times: float | Iterable[float]) -> np.ndarray:
        return calculate_spectrometers(self.liouvillian.l_free, times)

    def pulse_i(
        self,
        times: float | Iterable[float],
        phase: float,
        scale: float = 1.0,
    ) -> np.ndarray:
        dephased = self.b1_i_inh_scale == np.inf
        rad = phase * np.pi * 0.5
        liouv = (
            self.liouvillian.l_free
            + scale * np.cos(rad) * self.liouvillian.l_b1x_i
            + scale * np.sin(rad) * self.liouvillian.l_b1y_i
        )
        return calculate_spectrometers(liouv, times, dephased)

    def pulse_s(
        self,
        times: float | Iterable[float],
        phase: float,
        scale: float = 1.0,
    ) -> np.ndarray:
        rad = phase * np.pi * 0.5
        liouv = (
            self.liouvillian.l_free
            + scale * np.cos(rad) * self.liouvillian.l_b1x_s
            + scale * np.sin(rad) * self.liouvillian.l_b1y_s
        )
        return calculate_spectrometers(liouv, times)

    def pulse_is(
        self, times: float | Iterable[float], phase_i: float, phase_s: float
    ) -> np.ndarray:
        dephased = self.b1_i_inh_scale == np.inf
        liouv = (
            self.liouvillian.l_free
            + np.cos(phase_i * np.pi * 0.5) * self.liouvillian.l_b1x_i
            + np.sin(phase_i * np.pi * 0.5) * self.liouvillian.l_b1y_i
            + np.cos(phase_s * np.pi * 0.5) * self.liouvillian.l_b1x_s
            + np.sin(phase_s * np.pi * 0.5) * self.liouvillian.l_b1y_s
        )
        return calculate_spectrometers(liouv, times, dephased)

    def shaped_pulse_i(
        self, pw: float, amplitudes: Sequence[float], phases: Iterable[float]
    ) -> np.ndarray:
        time = pw / len(amplitudes)
        pairs = list(zip(amplitudes, phases))
        pulses = {
            (amp, ph): self.pulse_i(time, ph, scale=amp) for amp, ph in set(pairs)
        }
        return reduce(np.matmul, (pulses[pair] for pair in reversed(pairs)))

    def _calculate_base_pulses_i(self) -> None:
        if self.calculate_i_flag:
            pws = np.array([1.0, 2.0, 8.0 / 3.0]) * self._pw90_i
            base = self.pulse_i(pws, 0.0)
            pulses = self._add_phases(base, "i")
            self._p90_i, self._p180_i, self._p240_i = pulses.swapaxes(0, 1)

    @property
    def p90_i(self) -> np.ndarray:
        self._calculate_base_pulses_i()
        return self._p90_i

    @property
    def p180_i(self) -> np.ndarray:
        self._calculate_base_pulses_i()
        return self._p180_i

    @property
    def p240_i(self) -> np.ndarray:
        self._calculate_base_pulses_i()
        return self._p240_i

    @property
    def p9018090_i_1(self) -> np.ndarray:
        return self.p90_i[[3, 0, 1, 2]] @ self.p180_i @ self.p90_i[[3, 0, 1, 2]]

    @property
    def p9018090_i_2(self) -> np.ndarray:
        return self.p90_i[[1, 2, 3, 0]] @ self.p180_i @ self.p90_i[[1, 2, 3, 0]]

    @property
    def p9024090_i_1(self) -> np.ndarray:
        return self.p90_i[[3, 0, 1, 2]] @ self.p240_i @ self.p90_i[[3, 0, 1, 2]]

    @property
    def p9024090_i_2(self) -> np.ndarray:
        return self.p90_i[[1, 2, 3, 0]] @ self.p240_i @ self.p90_i[[1, 2, 3, 0]]

    def _calculate_base_pulses_s(self) -> None:
        if self.calculate_s_flag:
            pws = np.array([2.0]) * self._pw90_s
            base = self.pulse_s(pws, 0.0)
            pulses = self._add_phases(base, "s")
            self._p180_s = pulses.swapaxes(0, 1)

    @property
    def p180_s(self) -> np.ndarray:
        self._calculate_base_pulses_s()
        return self._p180_s

    def p9024090_nh(self, reverse: bool = False) -> np.ndarray:
        ph_n = 1 if reverse else 3
        ph_h = 3 if reverse else 1
        pw240i, pw9024090i = np.array([8.0, 14.0]) * self._pw90_i / 3.0
        pw240s, pw9024090s = np.array([8.0, 14.0]) * self._pw90_s / 3.0
        t0, t1, t2, t3 = 0.5 * np.diff(
            np.sort([pw240i, pw240s, pw9024090i, pw9024090s]), prepend=0.0
        )
        p0 = self.pulse_is(2.0 * t0, 0, 0)
        if pw9024090i <= pw9024090s:
            p1 = self.pulse_is(t1, ph_n, 0)
            p2 = (
                self.pulse_is(t2, ph_n, ph_h)
                if pw9024090i > pw240s
                else self.pulse_s(t2, 0)
            )
            p3 = self.pulse_s(t3, ph_h)
        else:
            p1 = self.pulse_is(t1, 0, ph_h)
            p2 = (
                self.pulse_is(t2, ph_n, ph_h)
                if pw9024090s > pw240i
                else self.pulse_i(t2, 0)
            )
            p3 = self.pulse_i(t3, ph_n)
        pw9024090is_xx = p3 @ p2 @ p1 @ p0 @ p1 @ p2 @ p3
        return self._add_phases(self._add_phases(pw9024090is_xx, "s"), "i")

    @property
    def p9024090_nh_1(self) -> np.ndarray:
        return self.p9024090_nh()

    @property
    def p9024090_nh_2(self) -> np.ndarray:
        return self.p9024090_nh(reverse=True)

    def calculate_shifts(self) -> np.ndarray:
        liouv = self.liouvillian.l_free.reshape(
            (self.liouvillian.size, self.liouvillian.size)
        )
        return eigvals(liouv).imag

    def offsets_to_ppms(self, offsets: np.ndarray) -> np.ndarray:
        return self.liouvillian.offsets_to_ppms(offsets)

    def ppms_to_offsets(self, ppms: np.ndarray | float) -> np.ndarray | float:
        return self.liouvillian.ppms_to_offsets(ppms)

max_stars_repo_name: gbouvignies/chemex

max_issues_repo_name: gbouvignies/chemex

max_forks_repo_name: gbouvignies/chemex


text:
import unittest

import numpy as np
import numpy.testing as testing

class basic(unittest.TestCase):
    def test_import(self):
        import idi.util  # noqa

    def test_angles(self):
        from idi.util import angles, rotation

        testing.assert_allclose(angles(np.eye(3)), (0, 0, 0), atol=1e-15)
        testing.assert_allclose(rotation(0, 0, 0), np.eye(3), atol=1e-15)
        for case in [(np.pi, 0, 0), (0, np.pi, 0), (0, 0, np.pi), (1, 1, 1)]:
            rm = rotation(*case)
            testing.assert_allclose(rotation(*angles(rm)), rm, atol=1e-15, err_msg=f"{case}: rotation(*angles({rm})) != {rm}")

    def test_fastlen(self):
        from idi.util import fastlen

        for x in [1, 2, 10, 17, 101, 5000]:
            self.assertGreaterEqual(fastlen(x), x)
            self.assertGreaterEqual(fastlen(x, factors=[2]), x)
            self.assertGreaterEqual(fastlen(x, factors=[2, 3]), x)

        xs = np.array([1, 11, 5001])
        for x, f in zip(xs, fastlen(xs, factors=[2, 3])):
            self.assertGreaterEqual(f, x)
        for x, f in zip(xs, fastlen(xs)):
            self.assertGreaterEqual(f, x)

        self.assertEqual(19, fastlen(19, [2, 19]))
        self.assertEqual(24, fastlen(19, [2, 3]))
        self.assertEqual(32, fastlen(19, [2]))
        self.assertEqual(6048, fastlen(6001, factors=[2, 3, 5, 7, 11]))

    def test_radial_profile(self):
        from idi.util import radial_profile

        # basic
        t = np.ones((100, 100))
        r, s = radial_profile(t, calcStd=True)
        testing.assert_allclose(r, 1, atol=1e-15)
        testing.assert_allclose(s, 0, atol=1e-15)
        self.assertEqual(len(r), 72)
        self.assertEqual(True, True)
        # center
        t = np.zeros((10, 10))
        t[2:5, :3] = 1
        r = radial_profile(t, center=(3, 1))
        testing.assert_allclose(r, np.array([1.0, 1.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0]))
        # random
        t = np.random.default_rng(0).random((50, 50))
        r, s = radial_profile(t, calcStd=True)
        testing.assert_array_less(r, 1)
        testing.assert_array_less(0, r)
        testing.assert_array_less(0, s[1:])
        testing.assert_allclose(r.mean(), 0.5, atol=0.1)
        # nan
        t[t > 0.8] = np.nan
        t[~np.isnan(t)] = 1
        r = radial_profile(t)
        testing.assert_allclose(r, 1, atol=1e-15)
        self.assertFalse(np.any(np.isnan(r)))

    def test_abs(self):
        from idi.util import abs2, abs2c

        self.assertAlmostEqual(abs2(1j + 1), 2)
        t = np.array(1 - 1j, dtype=np.complex128)
        self.assertAlmostEqual(abs2c(t), 2.0)
        self.assertFalse(np.iscomplex(abs2c(t)), 2.0)
        self.assertTrue(np.iscomplexobj(abs2c(t)), 2.0)

    def test_shortsci(self):
        from idi.util import shortsci

        for n, s in zip([1, 0.1, -0.1, 10, 1e17, -1e-17, 0.5, 1.5, -15], ["1e0", "1e-1", "-1e-1", "1e1", "1e17", "-1e-17", "5e-1", "2e0", "-2e1"]):
            self.assertEqual(shortsci(n), s)

        for n, s in zip([1, 0.555, 1.51, -1.234e17], ["1.0e0", "5.6e-1", "1.5e0", "-1.2e17"]):
            self.assertEqual(shortsci(n, decimals=1), s)

class random(unittest.TestCase):
    """
    TODO:
    poisson_disc_sample
    rndConvexpoly
    rndIcosahedron
    rndstr
    gnorm
    """

    def test_rndSphere(self):
        from idi.util import rndSphere

        R = 10
        N = int(1e6)
        t = rndSphere(R, N)
        r = np.linalg.norm(t, axis=1)
        self.assertTupleEqual(t.shape, (N, 3))
        self.assertGreater(r.max(), R * 0.95)
        self.assertLessEqual(r.max(), R)
        hist = np.histogram(r, bins=np.linspace(0, 1, int(N ** (1 / 3) * 0.2)) ** (1 / 3) * R)[0]
        self.assertLess(np.ptp(hist) / np.mean(hist), 0.05, msg="r deviation higher than 5%")
        d, b = np.histogramdd(t, density=True, bins=11)
        br = np.linalg.norm(np.meshgrid(*((c[:-1] + c[1:]) / 2 for c in b)), axis=0)
        f = d[br < 0.9 * R]
        self.assertLess(f.std() / f.mean(), 0.05, msg="density deviation higher than 5%")

    def test_rndstr(self):
        from idi.util import rndstr

        self.assertEqual(len(rndstr(0)), 0)
        self.assertEqual(len(rndstr(1)), 1)
        self.assertEqual(len(rndstr(1000)), 1000)
        s = set()
        for i in range(100):
            c = rndstr(5)
            s.add(c)
            self.assertTrue(c.isascii())
        self.assertEqual(len(s), 100, msg="repeats found")

    def test_rndgennorm(self):
        from idi.util import rndgennorm, fwhm

        fwhms = (1, 5, 10)
        t = rndgennorm(0, fwhms, (2, 2, 100), int(1e7))
        np.testing.assert_allclose(np.mean(t, axis=0), 0, atol=0.01, err_msg="mean failed")
        self.assertTupleEqual(t.shape, (int(1e7), 3), msg="shape missmatch")
        hist, bins = zip(*(np.histogram(x, bins=int(1e3), range=(-10, 10)) for x in t.T))
        for h, b, f in zip(hist, bins, fwhms):
            self.assertAlmostEqual(fwhm(b, h), f, delta=0.05, msg="fwhm failed")

    def test_randomrotation(self):
        from idi.util import random_rotation

        t = random_rotation()
        self.assertTupleEqual(t.shape, (3, 3))
        self.assertAlmostEqual(abs(np.linalg.det(t)), 1)

class h5util(unittest.TestCase):
    def setUp(self):
        import h5py

        self.f1 = h5py.File(name="f1", driver="core", backing_store=False, mode="w")
        self.f1.create_group("group")

    def tearDown(self):
        self.f1.close()

    def test_append(self):
        from idi.util import h5util

        h5util.appenddata(self.f1, "t1", np.zeros((1, 10, 10)))
        t = self.f1["t1"]
        self.assertTupleEqual(t.shape, (1, 10, 10))
        self.assertRaises(KeyError, h5util.appenddata, t, "blub", np.zeros((3, 10, 10)))
        h5util.appenddata(t, "/", np.zeros((3, 10, 10)))
        self.assertTupleEqual(t.shape, (4, 10, 10))
        self.assertEqual(t.compression, "lzf")
        self.assertEqual(t.shuffle, True)
        self.assertTupleEqual(t.chunks, (1, 10, 10))

        h5util.appenddata(self.f1, "string", "s1")
        h5util.appenddata(self.f1, "string", "s2_ismuchlonger")
        h5util.appenddata(self.f1, "string", np.array(self.f1["string"]))
        t = self.f1["string"]
        self.assertEqual(len(t), 4)
        self.assertEqual(t[-1], b"s2_ismuchlonger")
        self.assertEqual(np.array(t)[-2], b"s1")

        h5util.appenddata(self.f1["group"], "t2", np.ones(5))
        self.assertTupleEqual(self.f1["group/t2"].shape, (5,))

    def test_overwrite(self):
        from idi.util import h5util

        h5util.overwritedata(self.f1, "t1", np.zeros((1, 10, 10)))
        t = self.f1["t1"]
        self.assertTupleEqual(t.shape, (1, 10, 10))
        self.assertEqual(t.compression, "lzf")
        self.assertEqual(t.shuffle, True)
        self.assertTupleEqual(t.chunks, (1, 10, 10))

        h5util.overwritedata(self.f1, "t1", np.zeros((5, 10)), chunks=(2, 1))
        t = self.f1["t1"]
        self.assertTupleEqual(t.shape, (5, 10))
        self.assertTupleEqual(t.chunks, (2, 1))

        self.assertRaises(TypeError, h5util.overwritedata, t, "blub", np.zeros((3, 10, 10)))

        h5util.overwritedata(self.f1, "t2", "test")
        self.assertEqual(self.f1["t2"][0], b"test")

        h5util.overwritedata(self.f1, "t2", "othertest")
        self.assertEqual(self.f1["t2"][0], b"othertest")

        h5util.overwritedata(self.f1["group"], "t3", np.array(["a", "test"]))
        self.assertEqual(self.f1["group/t3"][1], b"test")

    def test_chunkediter(self):
        from idi.util import h5util

        self.f1["t4"] = np.zeros((15, 10, 1))

        for i, el in enumerate(h5util.chunkediter(self.f1["t4"], outsize=2)):
            if i < 7:
                self.assertTupleEqual(el.shape, (2, 10, 1), msg=f"differ at {i}: {el.shape}")
            else:
                self.assertTupleEqual(el.shape, (1, 10, 1), msg=f"differ at {i}: {el.shape}")
        self.assertEqual(i, 7)

        for i, el in enumerate(h5util.chunkediter(self.f1["t4"], outsize=1)):
            self.assertTupleEqual(el.shape, (10, 1), msg=f"differ at {i}: {el.shape}")
        self.assertEqual(i, 14)

        for i, el in enumerate(h5util.chunkediter(self.f1["t4"], outsize=5)):
            self.assertTupleEqual(el.shape, (5, 10, 1), msg=f"differ at {i}: {el.shape}")
        self.assertEqual(i, 2)

        # h5util.overwritedata(self.f1, 'string', 's1')
        # h5util.overwritedata(self.f1, 'string', 's2_ismuchlonger')
        # h5util.overwritedata(
        #     self.f1, 'string', np.array(self.f1['string']),
        # )
        # t = self.f1['string']
        # self.assertEqual(len(t), 4)
        # self.assertEqual(t[-1], b's2_ismuchlonger')
        # self.assertEqual(np.array(t)[-2], b's1')

class funchelper(unittest.TestCase):
    def test_asgen(self):
        from idi.util import asgen
        import types

        @asgen
        def test(el):
            return el

        t = test([1, 2, 3])
        self.assertIsInstance(t, types.GeneratorType)
        self.assertEqual(list(t), [1, 2, 3])

    def test_group(self):
        from idi.util import group

        a = range(10)
        for value, truth in zip(group(a, 4), [tuple(range(n, n + 4)) for n in range(0, 20, 4)]):
            self.assertTupleEqual(value, truth)

class array(unittest.TestCase):
    """
    TODO:
    create_mask
    diffdist
    fftfilter_mean
    fftfilter_std
    filter_std
    arrayfromiter
    split
    """

    def test_cutnan(self):
        from idi.util import cutnan

        t = np.zeros((10, 20))
        t[:5, :] = np.nan
        t[-5, :] = np.nan
        t[:, :2] = np.nan
        t[:, 15:17] = np.nan
        self.assertEqual(cutnan(t).shape, (4, 16))

        t = np.zeros((1, 20))
        t[:] = np.nan
        self.assertEqual(cutnan(t).shape, (0, 0))

    def test_centeredpart(self):
        from idi.util import centered_part

        t = np.arange(80).reshape(4, 4, 5)
        testing.assert_array_equal(centered_part(t, (1, 2, 2)), np.array([[[47, 48], [52, 53]]]))

    def test_list2array(self):
        from idi.util import list2array

        self.assertRaises(TypeError, list2array, 1)

        li = [1, 2, 3]
        testing.assert_array_equal(list2array(li), np.array([1, 2, 3]))

        li = [[1], [1, 2], [1, 2, 3], [1]]
        a = list2array(li)
        self.assertEqual(a[0, 0], li[0][0])
        self.assertEqual(a[1, 1], li[1][1])
        self.assertEqual(a[2, 2], li[2][2])
        self.assertEqual(np.sum(a), 11)
        self.assertEqual(len(a), len(li))

        li = [[1.0], ["a", "b"], ["a", "b", "c"], ["a"]]
        a = list2array(li)
        self.assertEqual(float(a[0, 0]), li[0][0])
        self.assertEqual(a[1, 1], li[1][1])
        self.assertEqual(a[2, 2], li[2][2])
        self.assertEqual(len(a), len(li))

    def test_fill(self):
        from idi.util import fill

        t = np.arange(20).reshape(4, 5).astype(float)
        t[1:4, 1:4] = 1
        t[2, 2] = np.nan
        t[0, 0] = np.nan
        t[-1, 0] = np.nan
        f = fill(t)
        good = np.copy(t)
        good[2, 2] = 1
        good[0, 0] = 5
        good[-1, 0] = 10
        self.assertFalse(np.any(np.isnan(f)))
        testing.assert_array_equal(good, f)

        t = np.empty((3, 3))
        t[:] = np.nan
        f = fill(t)
        self.assertTrue(np.all(np.isnan(f)))

        t = np.empty((3, 3))
        t[:] = np.nan
        t[1, 1] = 1
        f = fill(t)
        testing.assert_array_equal(f, 1)

    def test_atleastnd(self):
        from idi.util import atleastnd

        self.assertTupleEqual(atleastnd(0, 2).shape, (1, 1))
        self.assertTupleEqual(atleastnd(np.zeros((2, 2)), 1).shape, (2, 2))
        self.assertTupleEqual(atleastnd(np.zeros((3, 2)), 4).shape, (1, 1, 3, 2))

    def test_rebin(self):
        from idi.util import rebin

        t = np.arange(20).reshape(4, 5).astype(float)
        testing.assert_array_equal(rebin(t, (1, 1), "sum"), t)
        testing.assert_array_equal(rebin(t, (2, 2), "sum"), np.array([[12.0, 20.0], [52.0, 60.0]]))
        self.assertEqual(rebin(t, 4, "sum"), t[:, :4].sum())
        self.assertEqual(rebin(t, 3, "max"), t[:3, :3].max())
        self.assertEqual(rebin(t, 5, "max"), t.max())

        t = np.ones((4, 4), bool)
        self.assertEqual(rebin(t, 2, "min").dtype, bool)
        testing.assert_array_equal(
            rebin(t, 2,), np.ones((2, 2)),
        )

        t = np.ones((20, 2, 2))
        testing.assert_array_equal(rebin(t, (10, 20, 1), "mean"), np.ones((2, 1, 2)))

        self.assertRaises(ValueError, rebin, t, 1, "blub")
        self.assertRaises(ValueError, rebin, t, (1, 1), "sum")

    def test_indextreme(self):
        from idi.util import indmax, indmin

        t = np.array([1, 0, 1, 2, 1, np.nan])
        self.assertTupleEqual(indmax(t), (3,))
        self.assertTupleEqual(indmin(t), (1,))

        t = np.ones((4, 4, 4))
        t[1, 2, 3] = 2
        t[3, 2, 1:] = 0
        self.assertTupleEqual(indmax(t), (1, 2, 3))
        self.assertTupleEqual(indmin(t), (3, 2, 1))

    def test_intersect2d(self):
        from idi.util import intersect2d

        a = np.array([[1, 2, 3], [2, 3, 4], [1, 2, 3], [9, 9, 9]])
        b = np.array([[3, 2, 1], [4, 3, 2], [3, 2, 1], [9, 9, 9]])
        c = intersect2d(a, b)
        self.assertTupleEqual(c.shape, (1, 3))
        self.assertTupleEqual(tuple(c.ravel()), (9, 9, 9))

    def test_alignedarray(self):
        from idi.util import alignedarray

        t = alignedarray((3,), dtype=np.float32, alignment=64)
        self.assertEqual(t.shape, (3,))
        self.assertEqual(t.ctypes.data % 64, 0)

        t = alignedarray((2, 17), alignment=128, zero=True)
        self.assertEqual(t.shape, (2, 17))
        self.assertEqual(t.ctypes.data % 128, 0)

class accum(unittest.TestCase):
    def test_basic(self):
        from idi.util import accumulator

        a = accumulator(False)
        self.assertEqual(None, a.mean)
        self.assertEqual(None, a.std)
        self.assertEqual(a.n, 0)

        a.add(np.ones((2, 10, 10)))
        t = np.ones((2, 10, 10))
        t[0, ...] = 0
        a.add(t)
        self.assertEqual(a.n, 2)
        testing.assert_allclose(a.mean, np.array((0.5, 1))[:, None, None] * np.ones((2, 10, 10)))
        testing.assert_allclose(a.std, np.array((0.5, 0))[:, None, None] * np.ones((2, 10, 10)))
        testing.assert_allclose(a.result.mean, np.array((0.5, 1))[:, None, None] * np.ones((2, 10, 10)))
        testing.assert_allclose(a.result.std, np.array((0.5, 0))[:, None, None] * np.ones((2, 10, 10)))

    def test_shape(self):
        from idi.util import accumulator

        a = accumulator(False)
        a.add(np.ones((2, 10, 10)))
        self.assertTupleEqual(a.shape, (2, 10, 10))
        a.add(np.ones((1, 10, 1)))
        a.add(np.ones((1, 1, 1)))
        testing.assert_allclose(a.result.mean, np.ones((2, 10, 10)))
        self.assertRaises(ValueError, a.add, np.ones((3, 3, 3)))

    def test_minmax(self):
        from idi.util import accumulator

        a = accumulator(True)
        a.add(np.ones(3))
        a.add(np.array((0, 1, 2)))
        self.assertEqual(a.n, 2)
        testing.assert_allclose(a.result.mean, np.array((0.5, 1, 1.5)))
        testing.assert_allclose(a.result.std, np.array((0.5, 0, 0.5)))
        testing.assert_allclose(a.max, np.array((1, 1, 2)))
        testing.assert_allclose(a.min, np.array((0, 1, 1)))

    def test_combine(self):
        from idi.util import accumulator

        a = accumulator(True)
        a.add(np.ones(2))

        b = accumulator(False)
        b.add(np.ones(2))
        self.assertRaises(ValueError, a.combine, b)

        b = accumulator(True)
        b.add(np.ones(3))
        self.assertRaises(ValueError, a.combine, b)

        b = accumulator(True)
        b.add(np.zeros(2))
        a.combine(b)
        testing.assert_allclose(a.std, 0.5)
        testing.assert_allclose(a.mean, 0.5)
        testing.assert_allclose(a.max, 1)
        testing.assert_allclose(a.min, 0)
        self.assertEqual(a.n, 2)

if __name__ == "__main__":
    unittest.main()


text:
import numpy as np
from subprocess import call
import time
import pkg_resources
from sarabande.utils.functions import *

def calc_zeta(measure_obj, normalize=True, verbose_flag=True, skip_prepare=False):
    """
    This function is where the core algorithms take place for measuring the 3/4 PCFs
    either projected or not projected. In total there are 4 options

    - projected 3PCF        - projected 4PCF
    - full 3PCF             - full 4PCF

    Args:
        measure_obj (class: measure): an object that carries the necessary values and data structures to compute 3/4 PCFs
        normalize (bool, optional): flag to determine whether we would like to normalize by the bin volume(s). Defaults to True.
        verbose_flag (bool, optional): flag to walk the user through the full calculation process. Defaults to True.
        skip_prepare (bool, optional): Flag to determine whether or not we should skip preparing the data (if you already ran the calculation). Defaults to False.

    Raises:
        AssertionError: in order to calculate the full 4PCF one has to save a boundsandnumber file (this comes from creating radial bins in one of the utility functions.)

    Returns:
        No Returns, instead we add the zeta attribute to the measure_obj.
    """

    ##############################################################
    #                Projected Zeta Calculation
    ##############################################################
    if measure_obj.projected:
        start = time.time()
        if not skip_prepare:
            print("""Preparing the data:""")
            projected_prepare_data(measure_obj, verbose_flag)
        else:
            # fix bounds and number scope issue
            measure_obj.kernel_name = measure_obj.save_name
            # measure_obj.boundsandnumber = np.load(measure_obj.save_dir + 'bin_bounds_and_pixel_number_'+measure_obj.save_name+'.npy')

        if measure_obj.nPCF == 3:
            if verbose_flag:
                print("Starting Calculation of Projected 3PCF.")

            PCF_computed = (
                np.zeros((measure_obj.nbins, measure_obj.nbins, measure_obj.m_max)) + 0j
            )
            for bin1 in range(measure_obj.nbins):
                for bin2 in range(bin1):
                    for m in range(measure_obj.m_max):
                        PCF_computed[bin1, bin2, m] += np.sum(
                            measure_obj.density_field_data
                            * (
                                measure_obj.conv_M[m, bin1, :, :]
                                * measure_obj.conv_M[m, bin2, :, :].conjugate()
                            )
                        )

            PCF_computed_ = PCF_computed.copy()
            if normalize:
                # Normalize Full
                for bin1 in range(measure_obj.nbins):
                    for bin2 in range(bin1 + 1):
                        for m in range(measure_obj.m_max):
                            PCF_computed_[bin1, bin2, m] /= (
                                measure_obj.boundsandnumber[1, bin1]
                                * measure_obj.boundsandnumber[1, bin2]
                            )
                            PCF_computed_[bin2, bin1, m] = PCF_computed_[bin1, bin2, m]

                measure_obj.zeta = PCF_computed_
            else:
                measure_obj.zeta = PCF_computed_

            stop = time.time()

            print(
                "\nFinished Calculating the Projected 3PCF in {0:0.4f} seconds".format(
                    stop - start
                )
            )
            # coeff_1d = PCF_computed_1d_

        elif measure_obj.nPCF == 4:
            if verbose_flag:
                print("Starting Calculation of Projected 4PCF.")

            # compute the 4PCF
            PCF_computed = (
                np.zeros(
                    (
                        measure_obj.nbins,
                        measure_obj.nbins,
                        measure_obj.nbins,
                        2 * measure_obj.m_max,
                        2 * measure_obj.m_max,
                    )
                )
                + 0j
            )
            for bin1 in range(measure_obj.nbins):
                for bin2 in range(bin1):
                    for bin3 in range(bin2):
                        for m1 in range(-measure_obj.m_max, measure_obj.m_max):
                            if m1 > 0:
                                c_m1 = measure_obj.conv_M[m1, bin1, :, :]
                            else:
                                c_m1 = (-1) ** m1 * measure_obj.conv_M[
                                    -m1, bin1, :, :
                                ].conjugate()

                            for m2 in range(-measure_obj.m_max, measure_obj.m_max):
                                if m2 > 0:
                                    c_m2 = measure_obj.conv_M[m2, bin2, :, :]
                                else:
                                    c_m2 = (-1) ** m2 * measure_obj.conv_M[
                                        -m2, bin2, :, :
                                    ].conjugate()

                                m3 = -(m1 + m2)
                                if abs(m3) >= measure_obj.m_max:
                                    continue
                                else:
                                    if m3 > 0:
                                        c_m3 = measure_obj.conv_M[m3, bin3, :, :]
                                    else:
                                        c_m3 = (-1) ** m3 * measure_obj.conv_M[
                                            -m3, bin3, :, :
                                        ].conjugate()

                                    PCF_computed[bin1, bin2, bin3, m1, m2] += np.sum(
                                        measure_obj.density_field_data
                                        * (c_m1 * c_m2 * c_m3)
                                    )

            PCF_computed_ = PCF_computed.copy()
            if normalize:
                # Normalize and Fill
                for bin1 in range(measure_obj.nbins):
                    for bin2 in range(bin1):
                        for bin3 in range(bin2):
                            for m1 in range(-measure_obj.m_max, measure_obj.m_max):
                                for m2 in range(-measure_obj.m_max, measure_obj.m_max):
                                    m3 = -(m1 + m2)
                                    if abs(m3) >= measure_obj.m_max:
                                        continue
                                    else:
                                        PCF_computed_[bin1, bin2, bin3, m1, m2] /= (
                                            measure_obj.boundsandnumber[1, bin1]
                                            * measure_obj.boundsandnumber[1, bin2]
                                            * measure_obj.boundsandnumber[1, bin3]
                                        )
                                        PCF_computed_[
                                            bin3, bin1, bin2, m1, m2
                                        ] = PCF_computed_[bin1, bin2, bin3, m1, m2]
                                        PCF_computed_[
                                            bin2, bin3, bin1, m1, m2
                                        ] = PCF_computed_[bin1, bin2, bin3, m1, m2]
                                        PCF_computed_[
                                            bin1, bin3, bin2, m1, m2
                                        ] = PCF_computed_[bin1, bin2, bin3, m1, m2]
                                        PCF_computed_[
                                            bin2, bin1, bin3, m1, m2
                                        ] = PCF_computed_[bin1, bin2, bin3, m1, m2]
                                        PCF_computed_[
                                            bin3, bin2, bin1, m1, m2
                                        ] = PCF_computed_[bin1, bin2, bin3, m1, m2]

                measure_obj.zeta = PCF_computed_
            else:
                measure_obj.zeta = PCF_computed_

            stop = time.time()

            print(
                "\nFinished Calculating the Projected 4PCF in {0:0.4f} seconds".format(
                    stop - start
                )
            )

    ##############################################################
    #                   Full Zeta Calculation
    ##############################################################
    else:

        if not skip_prepare:
            print("""Preparing the data:""")
            prepare_data(measure_obj, verbose_flag)
        else:
            measure_obj.kernel_name = measure_obj.save_name
            # if not hasattr(measure_obj, 'boundsandnumber'):
            #     measure_obj.boundsandnumber = np.load(measure_obj.save_dir + 'bin_bounds_and_pixel_number_'+measure_obj.save_name+'.npy')

        if measure_obj.nPCF == 3:
            start = time.time()

            # load alms (results of convolution at fixed ell and for all m, to combine).
            zeta = (
                np.zeros(
                    (measure_obj.ell_max + 1, measure_obj.nbins, measure_obj.nbins)
                )
                + 0j
            )

            for l in range(0, measure_obj.ell_max + 1, 1):
                for bin1 in range(0, measure_obj.nbins, 1):
                    for bin2 in range(0, bin1, 1):
                        for m in range(0, l + 1, 1):
                            # load b1 and b2 at this m

                            ylm_b1 = np.load(
                                measure_obj.save_dir
                                + measure_obj.save_name
                                + "conv_data_kernel_"
                                + measure_obj.kernel_name
                                + "_"
                                + str(l)
                                + "_"
                                + str(m)
                                + "_bin_"
                                + str(bin1)
                                + ".npy"
                            ).astype(np.complex128)
                            ylm_b2 = np.load(
                                measure_obj.save_dir
                                + measure_obj.save_name
                                + "conv_data_kernel_"
                                + measure_obj.kernel_name
                                + "_"
                                + str(l)
                                + "_"
                                + str(m)
                                + "_bin_"
                                + str(bin2)
                                + ".npy"
                            ).astype(np.complex128)
                            # form half of sum at that m.
                            ylm_b1 *= ylm_b2.conjugate()
                            del ylm_b2

                            # average spatially now.
                            ylm_avg = np.sum(measure_obj.density_field_data * ylm_b1)
                            del ylm_b1
                            # add in negative m: so basically we are computing m's in sum term by
                            # term but with matching between + and - m.
                            if m > 0:
                                ylm_avg += ylm_avg.conjugate()
                            zeta[
                                l, bin1, bin2
                            ] += ylm_avg  # now accumulate so we get sum over m.

            zeta *= 1.0 / (4 * np.pi)

            for l in range(0, measure_obj.ell_max + 1, 1):
                for bin1 in range(0, measure_obj.nbins, 1):
                    for bin2 in range(
                        bin1, measure_obj.nbins, 1
                    ):  # ?? for bin2 in range(0, bin2<=bin1, 1):
                        zeta[l, bin1, bin2] = zeta[l, bin2, bin1]

            if normalize == True:
                binvolume = measure_obj.boundsandnumber[1, 0 : measure_obj.nbins]
                normedzeta = zeta / (binvolume[:, None] * binvolume[None, :])
                measure_obj.zeta = normedzeta

            else:
                measure_obj.zeta = zeta

            end = time.time()
            print("\n3PCF took {0:0.4f} seconds to finish".format(end - start))

            call(
                "rm "
                + measure_obj.save_dir
                + "bin_bounds_and_pixel_number_"
                + measure_obj.save_name
                + "*",
                shell=True,
            )
            call(
                "rm "
                + measure_obj.save_dir
                + measure_obj.save_name
                + "conv_data_kernel_"
                + measure_obj.kernel_name
                + "*",
                shell=True,
            )

        elif measure_obj.nPCF == 4:

            """
            This is the the big calculation for measuring the 4PCF coefficients.
            This code is adapted from Philcox et al. 2021 (Encore Paper)
            """

            def S(m):
                """
                A simple function defined in section 4.1.1 of
                Philcox et al. 2021 above algorithm 3. (Helper Function)
                """
                if m == 0:
                    return 1 / 2
                return 1

            def complex_modulus(z):
                """
                assume z to be an array of complex numbers of the form z = a + ib
                (Helper Function)
                """
                a = np.real(z)
                b = np.imag(z)

                return np.sqrt(a ** 2 + b ** 2)

            # we define these local variables out of laziness to change my original non-object oriented code
            ell_max = measure_obj.ell_max
            nbins = measure_obj.nbins
            if hasattr(measure_obj, "boundsandnumber") == False:
                raise AssertionError(
                    "You need to run measure_obj.create_radial_bins() first!"
                )

            binvolume = measure_obj.boundsandnumber[1, 0:nbins]

            start = time.time()
            print("Executing 4PCF Calculation ...")

            zeta = (
                np.zeros((ell_max + 1, ell_max + 1, ell_max + 1, nbins, nbins, nbins))
                + 0j
            )

            stream = pkg_resources.resource_stream(__name__, "utils/CG_Coeffs.npy")
            CG_Coefficients = np.load(stream)
            for l_1 in range(0, ell_max + 1):
                for l_2 in range(0, ell_max + 1):
                    for l_3 in range(np.abs(l_1 - l_2), min(l_1 + l_2, ell_max) + 1):
                        if (
                            l_1 + l_2 + l_3
                        ) % 2 != 0:  # we don't assume this to be true for Turbulent ISM
                            continue
                        for m_1 in range(-l_1, l_1 + 1):
                            for m_2 in range(-l_2, l_2 + 1):
                                m_3 = -m_1 - m_2
                                if m_3 > l_3 or m_3 < 0:
                                    continue
                                coupling_w = (
                                    measure_obj.density_field_data
                                    * (-1) ** (l_1 + l_2 + l_3)
                                    * CG_Coefficients[l_1, l_2, l_3, m_1, m_2, m_3]
                                )
                                for b_1 in range(
                                    0, nbins
                                ):  # might be an error with indexing here
                                    if m_1 < 0:
                                        a_lmb_1 = (-1) ** m_1 * (
                                            np.load(
                                                measure_obj.save_dir
                                                + measure_obj.save_name
                                                + "conv_data_kernel_"
                                                + measure_obj.kernel_name
                                                + "_"
                                                + str(l_1)
                                                + "_"
                                                + str(-m_1)
                                                + "_bin_"
                                                + str(b_1)
                                                + ".npy"
                                            ).astype(np.complex128)
                                        ).conjugate()
                                    #
                                    else:
                                        a_lmb_1 = np.load(
                                            measure_obj.save_dir
                                            + measure_obj.save_name
                                            + "conv_data_kernel_"
                                            + measure_obj.kernel_name
                                            + "_"
                                            + str(l_1)
                                            + "_"
                                            + str(m_1)
                                            + "_bin_"
                                            + str(b_1)
                                            + ".npy"
                                        ).astype(np.complex128)
                                    #
                                    for b_2 in range(
                                        b_1 + 1, nbins
                                    ):  # might be an error with indexing here
                                        if m_2 < 0:
                                            a_lmb_2 = (-1) ** m_2 * (
                                                np.load(
                                                    measure_obj.save_dir
                                                    + measure_obj.save_name
                                                    + "conv_data_kernel_"
                                                    + measure_obj.kernel_name
                                                    + "_"
                                                    + str(l_2)
                                                    + "_"
                                                    + str(-m_2)
                                                    + "_bin_"
                                                    + str(b_2)
                                                    + ".npy"
                                                ).astype(np.complex128)
                                            ).conjugate()
                                        #
                                        else:
                                            a_lmb_2 = np.load(
                                                measure_obj.save_dir
                                                + measure_obj.save_name
                                                + "conv_data_kernel_"
                                                + measure_obj.kernel_name
                                                + "_"
                                                + str(l_2)
                                                + "_"
                                                + str(m_2)
                                                + "_bin_"
                                                + str(b_2)
                                                + ".npy"
                                            ).astype(np.complex128)
                                        #
                                        for b_3 in range(
                                            b_2 + 1, nbins
                                        ):  # might be an error with indexing here
                                            a_lmb_3 = np.load(
                                                measure_obj.save_dir
                                                + measure_obj.save_name
                                                + "conv_data_kernel_"
                                                + measure_obj.kernel_name
                                                + "_"
                                                + str(l_3)
                                                + "_"
                                                + str(m_3)
                                                + "_bin_"
                                                + str(b_3)
                                                + ".npy"
                                            ).astype(np.complex128)
                                            #
                                            zeta[
                                                l_1, l_2, l_3, b_1, b_2, b_3
                                            ] += np.sum(
                                                2
                                                * S(m_3)
                                                * coupling_w
                                                * np.real(a_lmb_1 * a_lmb_2 * a_lmb_3)
                                            )

            for l1 in range(0, ell_max + 1):
                for l2 in range(0, ell_max + 1):
                    for l3 in range(np.abs(l1 - l2), min(l1 + l2, ell_max) + 1):
                        if (
                            l1 + l2 + l3
                        ) % 2 != 0:  # we don't assume this to be true for Turbulent ISM
                            continue
                        for b1 in range(0, nbins):
                            for b2 in range(b1 + 1, nbins):
                                for b3 in range(b2 + 1, nbins):
                                    this_4pcf = zeta[l1, l2, l3, b1, b2, b3]
                                    zeta[l3, l1, l2, b3, b1, b2] = this_4pcf
                                    zeta[l2, l3, l1, b2, b3, b1] = this_4pcf
                                    zeta[l1, l3, l2, b1, b3, b2] = this_4pcf
                                    zeta[l2, l1, l3, b2, b1, b3] = this_4pcf
                                    zeta[l3, l2, l1, b3, b2, b1] = this_4pcf

            finish = time.time()

            call(
                "rm "
                + measure_obj.save_dir
                + "bin_bounds_and_pixel_number_"
                + measure_obj.save_name
                + "*",
                shell=True,
            )
            call(
                "rm "
                + measure_obj.save_dir
                + measure_obj.save_name
                + "conv_data_kernel_"
                + measure_obj.kernel_name
                + "*",
                shell=True,
            )

            print(
                "Finished Calculating 4PCF in {0:0.4f} seconds".format(finish - start)
            )
            if normalize:
                """
                Normalize zeta^L_B (where L = {\ell_1, \ell_2, \ell_3} and B = {b_1, b_2, b_3}) hat
                coefficients from calc_zeta by dividing by bin volume
                """
                # binvolume = measure_obj.boundsandnumber[1,0:nbins]
                # normalize_coeff = (binvolume[:,None, None] * binvolume[None,:, None] * binvolume[None, None, :])
                normalize_coeff = (4.0 * np.pi) ** (3.0)
                zeta = normalize_coeff * zeta / ((measure_obj.ld_one_d ** 3))
                measure_obj.zeta = zeta
            else:
                print("your zeta coefficients are not properly ormalized.")
                measure_obj.zeta = zeta


text:
import numpy
import torch
import torch.nn.functional as F
from scipy.signal import convolve2d

from aydin.nn.pytorch.it_ptcnn import PTCNNImageTranslator
from aydin.nn.pytorch.models.psf_convolution import PSFConvolutionLayer
from aydin.util.log.log import lprint

def to_numpy(tensor):
    return tensor.clone().detach().cpu().numpy()

class PTCNNDeconvolution(PTCNNImageTranslator):
    """
    Pytorch-based CNN image deconvolution
    """

    def __init__(self, psf_kernel=None, broaden_psf=1, **kwargs):
        """
        Constructs a CNN image translator using the pytorch deep learning library.

        :param normaliser_type: normaliser type
        :param balance_training_data: balance data ? (limits number training entries per target value histogram bin)
        :param monitor: monitor to track progress of training externally (used by UI)
        """
        super().__init__(**kwargs)

        for i in range(broaden_psf):
            psf_kernel = numpy.pad(psf_kernel, (1,), mode='constant', constant_values=0)
            psf_kernel = convolve2d(
                psf_kernel,
                numpy.array([[0.5, 1, 0.5], [1, 1, 1], [0.5, 1, 0.5]]) / 7,
                'same',
            )

        psf_kernel /= psf_kernel.sum()
        psf_kernel = psf_kernel.astype(numpy.float32)

        self.psf_kernel = psf_kernel
        self.psf_kernel_tensor = torch.from_numpy(
            self.psf_kernel[numpy.newaxis, numpy.newaxis, ...]
        ).to(self.device)

        self.psfconv = PSFConvolutionLayer(self.psf_kernel).to(self.device)

        self.enforce_blind_spot = False

        self.sharpening = 0.01

    def _train_loop(self, data_loader, optimizer, loss_function):
        try:
            self.model.kernel_continuity_regularisation = False
        except AttributeError:
            lprint("Cannot deactivate kernel continuity regularisation")

        super()._train_loop(data_loader, optimizer, loss_function)

    def _additional_losses(self, translated_image, forward_model_image):

        # non-negativity loss:
        non_negativity_loss = F.relu(-translated_image)
        loss = non_negativity_loss.mean()

        # Sharpen loss_deconvolution:
        if self.sharpening and self.sharpening > 0:
            sum_values = torch.sum(translated_image, (2, 3), keepdim=True)
            entropy = self.sharpening * torch.sqrt(1 + translated_image / sum_values)
            entropy_loss_value = entropy.mean()
            loss += entropy_loss_value

        return loss

    def _forward_model(self, input):
        return self.psfconv(input)

max_stars_repo_name: royerloic/aydin

max_issues_repo_name: royerloic/aydin

max_forks_repo_name: royerloic/aydin


text:
# ---
# jupyter:
#   jupytext:
#     cell_metadata_json: true
#     comment_magics: false
#     formats: py:light,notebooks//ipynb
#     text_representation:
#       extension: .py
#       format_name: light
#       format_version: '1.5'
#       jupytext_version: 1.6.0
#   kernelspec:
#     display_name: Python 3
#     language: python
#     name: python3
# ---

# + [markdown] {"slideshow": {"slide_type": "slide"}}
# # Scipy
#
# ![scipy](images/scipyshiny_small.png)
#
# Scipy is the scientific Python ecosystem :
# - fft, linear algebra, scientific computation,...
# - scipy contains numpy, it can be considered as an extension of numpy.
# - the add-on toolkits [Scikits](https://scikits.appspot.com/scikits) complements scipy.

# + {"slideshow": {"slide_type": "skip"}, "execution": {"iopub.status.busy": "2020-09-12T14:17:15.936486Z", "iopub.execute_input": "2020-09-12T14:17:15.938394Z", "iopub.status.idle": "2020-09-12T14:17:16.429133Z", "shell.execute_reply": "2020-09-12T14:17:16.429710Z"}}
%matplotlib inline
%config InlineBackend.figure_format = 'retina'
import numpy as np
import scipy as sp
import matplotlib.pyplot as plt
plt.rcParams['figure.figsize'] = (10,6)

# + [markdown] {"slideshow": {"slide_type": "slide"}}
# ##  SciPy main packages
# - `constants` : Physical and mathematical constants
# - `fftpack` : Fast Fourier Transform routines
# - `integrate` : Integration and ordinary differential equation solvers
# - `interpolate` : Interpolation and smoothing splines
# - `io` : Input and Output
# - `linalg` : Linear algebra
# - `signal` : Signal processing
# - `sparse` : Sparse matrices and associated routines
#

# + {"slideshow": {"slide_type": "slide"}, "execution": {"iopub.status.busy": "2020-09-12T14:17:16.506164Z", "iopub.execute_input": "2020-09-12T14:17:16.507121Z", "iopub.status.idle": "2020-09-12T14:17:16.656919Z", "shell.execute_reply": "2020-09-12T14:17:16.657548Z"}}
from scipy.interpolate import interp1d
x = np.linspace(-1, 1, num=5)  # 5 points evenly spaced in [-1,1].
y = (x-1.)*(x-0.5)*(x+0.5)     # x and y are numpy arrays
f0 = interp1d(x,y, kind='zero')
f1 = interp1d(x,y, kind='linear')
f2 = interp1d(x,y, kind='quadratic')
f3 = interp1d(x,y, kind='cubic')
f4 = interp1d(x,y, kind='nearest')

# + {"slideshow": {"slide_type": "slide"}, "execution": {"iopub.status.busy": "2020-09-12T14:17:16.665449Z", "iopub.execute_input": "2020-09-12T14:17:16.698435Z", "iopub.status.idle": "2020-09-12T14:17:17.100571Z", "shell.execute_reply": "2020-09-12T14:17:17.101112Z"}}
xnew = np.linspace(-1, 1, num=40)
ynew = (xnew-1.)*(xnew-0.5)*(xnew+0.5)
plt.plot(x,y,'D',xnew,f0(xnew),':', xnew, f1(xnew),'-.',
                xnew,f2(xnew),'-',xnew ,f3(xnew),'--',
                xnew,f4(xnew),'--',xnew, ynew, linewidth=2)
plt.legend(['data','zero','linear','quadratic','cubic','nearest','exact'],
          loc='best');

# + {"slideshow": {"slide_type": "slide"}, "execution": {"iopub.status.busy": "2020-09-12T14:17:17.243719Z", "iopub.execute_input": "2020-09-12T14:17:17.244691Z", "iopub.status.idle": "2020-09-12T14:17:17.246852Z", "shell.execute_reply": "2020-09-12T14:17:17.247393Z"}}
from scipy.interpolate import interp2d
x,y = sp.mgrid[0:1:20j,0:1:20j]  #create the grid 20x20
z = np.cos(4*sp.pi*x)*np.sin(4*sp.pi*y) #initialize the field
T1=interp2d(x,y,z,kind='linear')
T2=interp2d(x,y,z,kind='cubic')
T3=interp2d(x,y,z,kind='quintic')

# + {"slideshow": {"slide_type": "slide"}, "execution": {"iopub.status.busy": "2020-09-12T14:17:17.255498Z", "iopub.execute_input": "2020-09-12T14:17:17.256326Z", "iopub.status.idle": "2020-09-12T14:17:17.974877Z", "shell.execute_reply": "2020-09-12T14:17:17.975517Z"}}
X,Y=sp.mgrid[0:1:100j,0:1:100j] #create the interpolation grid 100x100
# complex -> number of points, float -> step size
plt.figure(1)
plt.subplot(221) #Plot original data
plt.contourf(x,y,z)
plt.title('20x20')
plt.subplot(222) #Plot linear interpolation
plt.contourf(X,Y,T1(X[:,0],Y[0,:]))
plt.title('100x100 linear')
plt.subplot(223) #Plot cubic interpolation
plt.contourf(X,Y,T2(X[:,0],Y[0,:]))
plt.title('100x100 cubic')
plt.subplot(224) #Plot quintic interpolation
plt.contourf(X,Y,T3(X[:,0],Y[0,:]))
plt.title('100x100 quintic')

# + [markdown] {"slideshow": {"slide_type": "slide"}}
# ## FFT : scipy.fftpack
# - FFT dimension 1, 2 and n : fft, ifft (inverse), rfft (real), irfft, fft2 (dimension 2), ifft2, fftn (dimension n), ifftn.
# - Discrete cosinus transform : dct
# - Convolution product : convolve

# + {"slideshow": {"slide_type": "fragment"}, "execution": {"iopub.status.busy": "2020-09-12T14:17:17.982603Z", "iopub.execute_input": "2020-09-12T14:17:17.983607Z", "iopub.status.idle": "2020-09-12T14:17:22.623113Z", "shell.execute_reply": "2020-09-12T14:17:22.623685Z"}}
from numpy.fft import fft, ifft
x = np.random.random(2048)
%timeit ifft(fft(x))

# + {"slideshow": {"slide_type": "fragment"}, "execution": {"iopub.status.busy": "2020-09-12T14:17:22.628695Z", "iopub.execute_input": "2020-09-12T14:17:22.629465Z", "iopub.status.idle": "2020-09-12T14:17:26.723442Z", "shell.execute_reply": "2020-09-12T14:17:26.724035Z"}}
from scipy.fftpack import fft, ifft
x = np.random.random(2048)
%timeit ifft(fft(x))

# + [markdown] {"slideshow": {"slide_type": "slide"}}
# ## Linear algebra : scipy.linalg
# - Sovers, decompositions, eigen values. (same as numpy).
# - Matrix functions : expm, sinm, sinhm,...
# - Block matrices diagonal, triangular, periodic,...

# + {"slideshow": {"slide_type": "fragment"}, "execution": {"iopub.status.busy": "2020-09-12T14:17:26.732815Z", "iopub.execute_input": "2020-09-12T14:17:26.733727Z", "iopub.status.idle": "2020-09-12T14:17:26.737386Z", "shell.execute_reply": "2020-09-12T14:17:26.737956Z"}}
import scipy.linalg as spl
b=np.ones(5)
A=np.array([[1.,3.,0., 0.,0.],
           [ 2.,1.,-4, 0.,0.],
           [ 6.,1., 2,-3.,0.],
           [ 0.,1., 4.,-2.,-3.],
           [ 0.,0., 6.,-3., 2.]])
print("x=",spl.solve(A,b,sym_pos=False)) # LAPACK ( gesv ou posv )
AB=np.array([[0.,3.,-4.,-3.,-3.],
             [1.,1., 2.,-2., 2.],
             [2.,1., 4.,-3., 0.],
             [6.,1., 6., 0., 0.]])
print("x=",spl.solve_banded((2,1),AB,b)) # LAPACK ( gbsv )

# + {"slideshow": {"slide_type": "slide"}, "execution": {"iopub.status.busy": "2020-09-12T14:17:26.744189Z", "iopub.execute_input": "2020-09-12T14:17:26.746272Z", "shell.execute_reply": "2020-09-12T14:17:26.749814Z", "iopub.status.idle": "2020-09-12T14:17:26.750732Z"}}
P,L,U = spl.lu(A) #  P A = L U
np.set_printoptions(precision=3)
for M in (P,L,U):
    print(M, end="\n"+20*"-"+"\n")

# + [markdown] {"slideshow": {"slide_type": "slide"}}
# ##  CSC (Compressed Sparse Column)
#
# - All operations are optimized
# - Efficient "slicing" along axis=1.
# - Fast Matrix-vector product.
# - Conversion to other format could be costly.

# + {"slideshow": {"slide_type": "fragment"}, "execution": {"iopub.status.busy": "2020-09-12T14:17:26.809145Z", "iopub.execute_input": "2020-09-12T14:17:26.810468Z", "iopub.status.idle": "2020-09-12T14:17:26.814990Z", "shell.execute_reply": "2020-09-12T14:17:26.815744Z"}}
import scipy.sparse as spsp
row = np.array([0,2,2,0,1,2])
col = np.array([0,0,1,2,2,2])
data  = np.array([1,2,3,4,5,6])
Mcsc1 = spsp.csc_matrix((data,(row,col)),shape=(3,3))
Mcsc1.todense()

# + {"slideshow": {"slide_type": "fragment"}, "execution": {"iopub.status.busy": "2020-09-12T14:17:26.826295Z", "iopub.execute_input": "2020-09-12T14:17:26.827566Z", "iopub.status.idle": "2020-09-12T14:17:26.832313Z", "shell.execute_reply": "2020-09-12T14:17:26.833136Z"}}
indptr  = np.array([0,2,3,6])
indices = np.array([0,2,2,0,1,2])
data    = np.array([1,2,3,4,5,6])
Mcsc2 = spsp.csc_matrix ((data,indices,indptr),shape=(3,3))
Mcsc2.todense()

# + [markdown] {"slideshow": {"slide_type": "slide"}}
# ## Dedicated format for assembling
# - `lil_matrix` : Row-based linked list matrix. Easy format to build your matrix and convert to other format before solving.
# - `dok_matrix` : A dictionary of keys based matrix. Ideal format for
# incremental matrix building. The conversion to csc/csr format is efficient.
# - `coo_matrix`  : coordinate list format. Fast conversion to formats CSC/CSR.
#
# [Lien vers la documentation scipy](http://docs.scipy.org/doc/scipy/reference/sparse.html)

# + [markdown] {"slideshow": {"slide_type": "slide"}}
# ## Sparse matrices : [scipy.sparse.linalg](https://docs.scipy.org/doc/scipy/reference/generated/scipy.sparse.linalg)
#
# - speigen, speigen_symmetric, lobpcg : (ARPACK).
# - svd : (ARPACK).
# - Direct methods (UMFPACK or SUPERLU) or Krylov based methods
# - Minimization : lsqr and minres
#
# For linear algebra:
# - Noobs: spsolve.
# - Intermmediate: dsolve.spsolve or isolve.spsolve
# - Advanced: splu, spilu (direct); cg, cgs, bicg, bicgstab, gmres, lgmres et qmr (iterative)
# - Boss: petsc4py et slepc4py.
#

# + [markdown] {"slideshow": {"slide_type": "slide"}}
# ## LinearOperator
#
# The LinearOperator is used for matrix-free numerical methods.

# + {"slideshow": {"slide_type": "fragment"}, "execution": {"iopub.status.busy": "2020-09-12T14:17:26.840147Z", "iopub.execute_input": "2020-09-12T14:17:26.841279Z", "iopub.status.idle": "2020-09-12T14:17:26.875505Z", "shell.execute_reply": "2020-09-12T14:17:26.876473Z"}}
import scipy.sparse.linalg as spspl
def mv(v):
   return np.array([2*v[0],3*v[1]])

A=spspl.LinearOperator((2 ,2),matvec=mv,dtype=float )
A

# + {"slideshow": {"slide_type": "slide"}, "execution": {"iopub.status.busy": "2020-09-12T14:17:26.882818Z", "iopub.execute_input": "2020-09-12T14:17:26.884155Z", "iopub.status.idle": "2020-09-12T14:17:26.888584Z", "shell.execute_reply": "2020-09-12T14:17:26.889395Z"}}
A*np.ones(2)

# + {"slideshow": {"slide_type": "fragment"}, "execution": {"iopub.status.busy": "2020-09-12T14:17:26.896480Z", "iopub.execute_input": "2020-09-12T14:17:26.897764Z", "iopub.status.idle": "2020-09-12T14:17:26.902265Z", "shell.execute_reply": "2020-09-12T14:17:26.903078Z"}}
A.matmat(np.array([[1,-2],[3,6]]))

# + [markdown] {"slideshow": {"slide_type": "slide"}}
# ## LU decomposition

# + {"slideshow": {"slide_type": "fragment"}, "execution": {"iopub.status.busy": "2020-09-12T14:17:26.912563Z", "iopub.execute_input": "2020-09-12T14:17:26.913782Z", "iopub.status.idle": "2020-09-12T14:17:26.920074Z", "shell.execute_reply": "2020-09-12T14:17:26.920886Z"}}
N = 50
un = np.ones(N)
w = np.random.rand(N+1)
A = spsp.spdiags([w[1:],-2*un,w[:-1]],[-1,0,1],N,N) # tridiagonal matrix
A = A.tocsc()
b = un
op = spspl.splu(A)
op

# + {"slideshow": {"slide_type": "fragment"}, "execution": {"iopub.status.busy": "2020-09-12T14:17:26.926070Z", "iopub.execute_input": "2020-09-12T14:17:26.927415Z", "iopub.status.idle": "2020-09-12T14:17:26.931809Z", "shell.execute_reply": "2020-09-12T14:17:26.932624Z"}}
x=op.solve(b)
spl.norm(A*x-b)

# + [markdown] {"slideshow": {"slide_type": "slide"}}
# ## Conjugate Gradient

# + {"slideshow": {"slide_type": "fragment"}, "execution": {"iopub.status.busy": "2020-09-12T14:17:26.941001Z", "iopub.execute_input": "2020-09-12T14:17:26.946217Z", "iopub.status.idle": "2020-09-12T14:17:26.949574Z", "shell.execute_reply": "2020-09-12T14:17:26.950147Z"}}
global k
k=0
def f(xk): # function called at every iterations
     global k
     print ("iteration {0:2d} residu = {1:7.3g}".format(k,spl.norm(A*xk-b)))
     k += 1

x,info=spspl.cg(A,b,x0=np.zeros(N),tol=1.0e-12,maxiter=N,M=None,callback=f)

# + [markdown] {"slideshow": {"slide_type": "slide"}}
# ## Preconditioned conjugate gradient

# + {"slideshow": {"slide_type": "fragment"}, "execution": {"iopub.status.busy": "2020-09-12T14:17:26.954769Z", "iopub.execute_input": "2020-09-12T14:17:26.955871Z", "iopub.status.idle": "2020-09-12T14:17:26.959314Z", "shell.execute_reply": "2020-09-12T14:17:26.959905Z"}}
pc=spspl.spilu(A,drop_tol=0.1)  # pc is an ILU decomposition
xp=pc.solve(b)
spl.norm(A*xp-b)

# + {"slideshow": {"slide_type": "fragment"}, "execution": {"iopub.status.busy": "2020-09-12T14:17:26.966049Z", "iopub.execute_input": "2020-09-12T14:17:26.966944Z", "iopub.status.idle": "2020-09-12T14:17:26.973399Z", "shell.execute_reply": "2020-09-12T14:17:26.973984Z"}}
def mv(v):
    return pc.solve(v)
lo = spspl.LinearOperator((N,N),matvec=mv)
k = 0
x,info=spspl.cg(A,b,x0=np.zeros(N),tol=1.e-12,maxiter=N,M=lo,callback=f)

# + [markdown] {"slideshow": {"slide_type": "slide"}}
# ## Numerical integration
#
# - quad, dblquad, tplquad,... Fortran library QUADPACK.
#

# + {"slideshow": {"slide_type": "fragment"}, "execution": {"iopub.status.busy": "2020-09-12T14:17:26.978341Z", "iopub.execute_input": "2020-09-12T14:17:26.979147Z", "iopub.status.idle": "2020-09-12T14:17:27.082255Z", "shell.execute_reply": "2020-09-12T14:17:27.082829Z"}}
import scipy.integrate as spi

x2=lambda x: x**2
4.**3/3  # int(x2) in [0,4]

# + {"slideshow": {"slide_type": "fragment"}, "execution": {"iopub.status.busy": "2020-09-12T14:17:27.087471Z", "iopub.execute_input": "2020-09-12T14:17:27.088374Z", "iopub.status.idle": "2020-09-12T14:17:27.090667Z", "shell.execute_reply": "2020-09-12T14:17:27.091224Z"}}
spi.quad(x2,0.,4.)

# + [markdown] {"slideshow": {"slide_type": "slide"}}
# ## Scipy ODE solver
#
# It uses the Fortran ODEPACK library.
#
# ### Van der Pol Oscillator
# \begin{align}
# y_1'(t)	& = y_2(t), \nonumber \\
# y_2'(t)	& = 1000(1 - y_1^2(t))y_2(t) - y_1(t) \nonumber
# \end{align}
# $$
# with $y_1(0) = 2 $ and $ y_2(0) = 0. $.

# + {"slideshow": {"slide_type": "fragment"}, "execution": {"iopub.status.busy": "2020-09-12T14:17:27.096802Z", "iopub.execute_input": "2020-09-12T14:17:27.097666Z", "iopub.status.idle": "2020-09-12T14:17:27.099006Z", "shell.execute_reply": "2020-09-12T14:17:27.099559Z"}}
import numpy as np
import scipy.integrate as spi

def vdp1000(y,t):
     dy=np.zeros(2)
     dy[0]=y[1]
     dy[1]=1000.*(1.-y[0]**2)*y[1]-y[0]
     return dy

# + {"slideshow": {"slide_type": "slide"}, "execution": {"iopub.status.busy": "2020-09-12T14:17:27.103955Z", "iopub.execute_input": "2020-09-12T14:17:27.104850Z", "iopub.status.idle": "2020-09-12T14:17:27.107325Z", "shell.execute_reply": "2020-09-12T14:17:27.107893Z"}}
t0, tf =0,  3000
N = 300000
t, dt = np.linspace(t0,tf,N, retstep=True)

# + {"slideshow": {"slide_type": "fragment"}, "execution": {"iopub.status.busy": "2020-09-12T14:17:27.112488Z", "iopub.execute_input": "2020-09-12T14:17:27.113396Z", "iopub.status.idle": "2020-09-12T14:17:27.424693Z", "shell.execute_reply": "2020-09-12T14:17:27.425407Z"}}
y=spi.odeint(vdp1000,[2.,0.],t)
plt.plot(t,y[:,0]);
# -

# ## Exercise
#
# The following code solve the Laplace equation using a dense matrix.
# - Modified the code to use a sparse matrix

# + {"execution": {"iopub.status.busy": "2020-09-12T14:17:27.462080Z", "iopub.execute_input": "2020-09-12T14:17:27.462973Z", "iopub.status.idle": "2020-09-12T14:17:28.327353Z", "shell.execute_reply": "2020-09-12T14:17:28.328125Z"}}
%matplotlib inline
%config InlineBackend.figure_format = "retina"
import numpy as np
import matplotlib.pyplot as plt
plt.rcParams['figure.figsize'] = (10,6)

# Boundary conditions
Tnorth, Tsouth, Twest, Teast = 100, 20, 50, 50

# Set meshgrid
n = 50
l = 1.0
h = l / (n-1)
X, Y = np.meshgrid(np.linspace(0,l,n), np.linspace(0,l,n))
T = np.zeros((n,n),dtype='d')

# Set Boundary condition
T[n-1:, :] = Tnorth / h**2
T[:1, :] = Tsouth / h**2
T[:, n-1:] = Teast / h**2
T[:, :1] = Twest / h**2

A = np.zeros((n*n,n*n),dtype='d')
nn = n*n
ii = 0
for j in range(n):
    for i in range(n):
      if j > 0:
         jj = ii - n
         A[ii,jj] = -1
      if j < n-1:
         jj = ii + n
         A[ii,jj] = -1
      if i > 0:
         jj = ii - 1
         A[ii,jj] = -1
      if i < n-1:
         jj = ii + 1
         A[ii,jj] = -1
      A[ii,ii] = 4
      ii = ii+1

# + {"execution": {"iopub.status.busy": "2020-09-12T14:17:27.462080Z", "iopub.execute_input": "2020-09-12T14:17:27.462973Z", "iopub.status.idle": "2020-09-12T14:17:28.327353Z", "shell.execute_reply": "2020-09-12T14:17:28.328125Z"}}
%%time
U = np.linalg.solve(A,np.ravel(h**2*T))

# + {"execution": {"iopub.status.busy": "2020-09-12T14:17:27.462080Z", "iopub.execute_input": "2020-09-12T14:17:27.462973Z", "iopub.status.idle": "2020-09-12T14:17:28.327353Z", "shell.execute_reply": "2020-09-12T14:17:28.328125Z"}}
T = U.reshape(n,n)
plt.contourf(X,Y,T)
plt.colorbar()

# + {"execution": {"iopub.status.busy": "2020-09-12T14:17:28.336444Z", "iopub.execute_input": "2020-09-12T14:17:28.337307Z", "iopub.status.idle": "2020-09-12T14:17:29.339935Z", "shell.execute_reply": "2020-09-12T14:17:29.340588Z"}}
import scipy.sparse as spsp
import scipy.sparse.linalg as spspl

# Boundary conditions
Tnorth, Tsouth, Twest, Teast = 100, 20, 50, 50

# Set meshgrid
n = 50
l = 1.0
h = l / (n-1)
X, Y = np.meshgrid(np.linspace(0,l,n), np.linspace(0,l,n))
T = np.zeros((n,n),dtype='d')

# Set Boundary condition
T[n-1:, :]    = Tnorth / h**2
T[  :1, :]    = Tsouth / h**2
T[   :, n-1:] = Teast  / h**2
T[   :, :1]   = Twest  / h**2

bdiag = -4 * np.eye(n)
bup   = np.diag([1] * (n - 1), 1)
blow  = np.diag([1] * (n - 1), -1)
block = bdiag + bup + blow
# Creat a list of n blocks
blist = [block] * n
S = spsp.block_diag(blist)
# Upper diagonal array offset by -n
upper = np.diag(np.ones(n * (n - 1)), n)
# Lower diagonal array offset by -n
lower = np.diag(np.ones(n * (n - 1)), -n)
S += upper + lower

# + {"execution": {"iopub.status.busy": "2020-09-12T14:17:28.336444Z", "iopub.execute_input": "2020-09-12T14:17:28.337307Z", "iopub.status.idle": "2020-09-12T14:17:29.339935Z", "shell.execute_reply": "2020-09-12T14:17:29.340588Z"}}
%%time
U = sp.linalg.solve(S,np.ravel(h**2*T))

# + {"execution": {"iopub.status.busy": "2020-09-12T14:17:28.336444Z", "iopub.execute_input": "2020-09-12T14:17:28.337307Z", "iopub.status.idle": "2020-09-12T14:17:29.339935Z", "shell.execute_reply": "2020-09-12T14:17:29.340588Z"}}
plt.contourf(X,Y,U.reshape(n,n))
plt.colorbar();

max_stars_repo_name: zhongyangynag/code-study

max_issues_repo_name: zhongyangynag/code-study

max_forks_repo_name: zhongyangynag/code-study

reason: import numpy,from numpy,import scipy,from scipy


text:
"""
Copyright (c) 2018-2022 Intel Corporation

Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at

      http://www.apache.org/licenses/LICENSE-2.0

Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
"""

from enum import Enum
import numpy as np

from .base_representation import BaseRepresentation
from ..data_readers import BaseReader

class GTLoader(Enum):
    PILLOW = 0
    OPENCV = 1

class StyleTransferRepresentation(BaseRepresentation):
    pass

class StyleTransferAnnotation(StyleTransferRepresentation):
    LOADERS = {
        GTLoader.PILLOW: 'pillow_imread',
        GTLoader.OPENCV: 'opencv_imread'
    }

    def __init__(self, identifier, image_path, gt_loader=GTLoader.PILLOW):
        super().__init__(identifier)
        self._gt_loader = self.LOADERS.get(gt_loader)
        self._image_path = image_path
        self._value = None

    @property
    def value(self):
        if self._value is not None:
            return self._value

        loader = BaseReader.provide(self._gt_loader, self.metadata['data_source'])
        gt = loader.read(self._image_path)
        return gt.astype(np.uint8)

    @value.setter
    def value(self, val):
        self._value = val
        return self._value

class StyleTransferPrediction(StyleTransferRepresentation):
    def __init__(self, identifiers, prediction):
        super().__init__(identifiers)
        self.value = prediction