text:
# Copyright 2018 Amazon.com, Inc. or its affiliates. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License").
# You may not use this file except in compliance with the License.
# A copy of the License is located at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# or in the "license" file accompanying this file. This file 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.

# Standard library imports
from typing import List

# Third-party imports
import mxnet as mx
import numpy as np

# First-party imports
from gluonts.core.component import validated
from gluonts.evaluation import get_seasonality
from gluonts.model.common import Tensor

VALID_N_BEATS_STACK_TYPES = "G", "S", "T"
VALID_LOSS_FUNCTIONS = "sMAPE", "MASE", "MAPE"

def linear_space(
    F, backcast_length: int, forecast_length: int, fwd_looking: bool
):
    if fwd_looking:
        return F.arange(0, forecast_length) / forecast_length
    else:
        return F.arange(-backcast_length, 0) / backcast_length  # Option 01
        # return F.arange(backcast_length, 0, -1) / backcast_length # Option 02
        # return F.arange(0, backcast_length) / backcast_length # Option 03

def seasonality_model(
    F,
    num_coefficients: int,
    context_length: int,
    prediction_length: int,
    is_forecast: bool,
) -> Tensor:
    """
    Creates a fourier series basis with num_coefficients coefficients for sine and cosine each.
    So the total number of learned coefficients amounts to 2*num_coefficients.
    """
    t = linear_space(
        F, context_length, prediction_length, fwd_looking=is_forecast
    )
    cosines = F.stack(
        *[F.cos(2 * np.pi * i * t) for i in range(num_coefficients)]
    )
    sines = F.stack(
        *[F.sin(2 * np.pi * i * t) for i in range(num_coefficients)]
    )
    S = F.concat(cosines, sines, dim=0)
    return S

def trend_model(
    F,
    num_coefficients: int,
    context_length: int,
    prediction_length: int,
    is_forecast: bool,
) -> Tensor:
    """
    Creates a polynomial basis of degree num_coefficients-1.
    """
    t = linear_space(
        F, context_length, prediction_length, fwd_looking=is_forecast
    )
    T = F.stack(*[t ** i for i in range(num_coefficients)])
    return T

class NBEATSBlock(mx.gluon.HybridBlock):
    """
    The NBEATS Block as described in the paper: https://arxiv.org/abs/1905.10437.

    Parameters
    ----------
    width
        Width of the fully connected layers with ReLu activation.
    num_block_layers
        Number of fully connected layers with ReLu activation.
    expansion_coefficient_length
        If the type is "G" (generic), then the length of the expansion coefficient.
        If type is "T" (trend), then it corresponds to the degree of the polynomial.
        If the type is "S" (seasonal) then its not used.
    prediction_length
        Length of the prediction. Also known as 'horizon'.
    context_length
        Number of time units that condition the predictions
        Also known as 'lookback period'.
    has_backcast
        Only the last block of the network doesn't.
    kwargs
        Arguments passed to 'HybridBlock'.
    """

    # Needs the validated decorator so that arguments types are checked and
    # the block can be serialized.
    @validated()
    def __init__(
        self,
        width: int,
        num_block_layers: int,
        expansion_coefficient_length: int,
        prediction_length: int,
        context_length: int,
        has_backcast: bool,
        **kwargs,
    ) -> None:
        super().__init__(**kwargs)

        self.width = width
        self.num_block_layers = num_block_layers
        self.expansion_coefficient_length = expansion_coefficient_length
        self.prediction_length = prediction_length
        self.context_length = context_length
        self.has_backcast = has_backcast

        # Parameter to control whether the basis matrix has
        # been initialised, if the block uses one
        self.basis_initialized = False

        # Only fc_stack defined concretely in this class
        with self.name_scope():
            self.fc_stack = mx.gluon.nn.HybridSequential()
            for i in range(self.num_block_layers):
                self.fc_stack.add(
                    mx.gluon.nn.Dense(
                        units=self.width,
                        activation="relu",
                        prefix=f"fc_stack_dense_{i}_",
                    )
                )

            # Subclasses will have to initialize these attributes appropriately:

            self.theta_backcast = None
            self.theta_forecast = None

            self.backcast = None
            self.forecast = None

            self.backcast_basis = None
            self.forecast_basis = None

    # This function is called upon first call of the hybrid_forward method
    def initialize_basis(self, F):
        pass

    def hybrid_forward(self, F, x, *args, **kwargs):
        # We do this to cache the constant basis matrix between forward passes
        if not self.basis_initialized:
            self.initialize_basis(F)
            self.basis_initialized = True

        x = self.fc_stack(x)
        theta_f = self.theta_forecast(x)
        forecast = self.forecast(theta_f)

        if self.has_backcast:
            theta_b = self.theta_backcast(x)
            backcast = self.backcast(theta_b)
            return backcast, forecast

        return forecast

class NBEATSGenericBlock(NBEATSBlock):
    """
    The NBEATS Block as described in the paper: https://arxiv.org/abs/1905.10437.
    This is the GenericBlock variant.

    Parameters
    ----------
    width
        Width of the fully connected layers with ReLu activation.
    num_block_layers
        Number of fully connected layers with ReLu activation.
    expansion_coefficient_length
        The length of the expansion coefficient.
    prediction_length
        Length of the prediction. Also known as 'horizon'.
    context_length
        Number of time units that condition the predictions
        Also known as 'lookback period'.
    has_backcast
        Only the last block of the network doesn't.
    kwargs
        Arguments passed to 'HybridBlock'.
    """

    # Needs the validated decorator so that arguments types are checked and
    # the block can be serialized.
    @validated()
    def __init__(self, **kwargs) -> None:
        super().__init__(**kwargs)

        with self.name_scope():
            if self.has_backcast:
                self.theta_backcast = mx.gluon.nn.Dense(
                    units=self.expansion_coefficient_length,
                    prefix=f"theta_backcast_dense_",  # linear activation:
                )
                self.backcast = mx.gluon.nn.Dense(
                    units=self.context_length,
                    prefix=f"backcast_dense_",  # linear activation:
                )
            self.theta_forecast = mx.gluon.nn.Dense(
                units=self.expansion_coefficient_length,
                prefix=f"theta_forecast_dense_",  # linear activation:
            )
            self.forecast = mx.gluon.nn.Dense(
                units=self.prediction_length,
                prefix=f"theta_dense_",  # linear activation:
            )

class NBEATSSeasonalBlock(NBEATSBlock):
    """
    The NBEATS Block as described in the paper: https://arxiv.org/abs/1905.10437.
    This is the Seasonal block variant.

    Parameters
    ----------
    width
        Width of the fully connected layers with ReLu activation.
    num_block_layers
        Number of fully connected layers with ReLu activation.
    expansion_coefficient_length
        Not used in this block type.
    prediction_length
        Length of the prediction. Also known as 'horizon'.
    context_length
        Number of time units that condition the predictions
        Also known as 'lookback period'.
    has_backcast
        Only the last block of the network doesn't.
    kwargs
        Arguments passed to 'HybridBlock'.
    """

    # Needs the validated decorator so that arguments types are checked and
    # the block can be serialized.
    @validated()
    def __init__(self, **kwargs) -> None:
        super().__init__(**kwargs)

        # the number of coefficient in the fourier basis per sine and cosine each
        # determined depending on prediction length, as defined by paper
        # also: dont use floor because prediction_length=1 should be mapped to 0
        self.num_coefficients = int((self.prediction_length / 2) - 1) + 1

        with self.name_scope():
            if self.has_backcast:
                self.theta_backcast = mx.gluon.nn.Dense(
                    units=2 * self.num_coefficients,
                    prefix=f"theta_backcast_dense_",  # linear activation:
                )
                self.backcast = mx.gluon.nn.HybridLambda(
                    lambda F, thetas: F.dot(thetas, self.backcast_basis),
                    prefix=f"backcast_lambda_",
                )
            self.theta_forecast = mx.gluon.nn.Dense(
                units=2 * self.num_coefficients,
                prefix=f"theta_forecast_dense_",  # linear activation:
            )
            self.forecast = mx.gluon.nn.HybridLambda(
                lambda F, thetas: F.dot(thetas, self.forecast_basis),
                prefix=f"forecast_lambda_",
            )

    def initialize_basis(self, F):
        # these are essentially constant matrices of type F that
        # define the basis for the seasonal model
        if self.has_backcast:
            self.backcast_basis = seasonality_model(
                F,
                num_coefficients=self.num_coefficients,
                context_length=self.context_length,
                prediction_length=self.prediction_length,
                is_forecast=False,
            )
        self.forecast_basis = seasonality_model(
            F,
            num_coefficients=self.num_coefficients,
            context_length=self.context_length,
            prediction_length=self.prediction_length,
            is_forecast=True,
        )

class NBEATSTrendBlock(NBEATSBlock):
    """"
    The NBEATS Block as described in the paper: https://arxiv.org/abs/1905.10437.
    This is the Trend block variant.

    Parameters
    ----------
    width
        Width of the fully connected layers with ReLu activation.
    num_block_layers
        Number of fully connected layers with ReLu activation.
    expansion_coefficient_length
        The length of the number of expansion coefficients.
        This corresponds to degree of the polynomial basis as follows:
            expansion_coefficient_length-1
    prediction_length
        Length of the prediction. Also known as 'horizon'.
    context_length
        Number of time units that condition the predictions
        Also known as 'lookback period'.
    has_backcast
        Only the last block of the network doesn't.
    kwargs
        Arguments passed to 'HybridBlock'.
    """

    # Needs the validated decorator so that arguments types are checked and
    # the block can be serialized.
    @validated()
    def __init__(self, **kwargs) -> None:
        super().__init__(**kwargs)

        with self.name_scope():
            if self.has_backcast:
                self.theta_backcast = mx.gluon.nn.Dense(
                    units=self.expansion_coefficient_length,
                    prefix=f"theta_backcast_dense_",  # linear activation:
                )
                self.backcast = mx.gluon.nn.HybridLambda(
                    lambda F, thetas: F.dot(thetas, self.backcast_basis),
                    prefix=f"backcast_lambda_",
                )
            self.theta_forecast = mx.gluon.nn.Dense(
                units=self.expansion_coefficient_length,
                prefix=f"theta_forecast_dense_",  # linear activation:
            )
            self.forecast = mx.gluon.nn.HybridLambda(
                lambda F, thetas: F.dot(thetas, self.forecast_basis),
                prefix=f"forecast_lambda_",
            )

    def initialize_basis(self, F):
        # these are essentially constant matrices of type F that
        # define the basis for the trend model
        if self.has_backcast:
            self.backcast_basis = trend_model(
                F,
                num_coefficients=self.expansion_coefficient_length,
                context_length=self.context_length,
                prediction_length=self.prediction_length,
                is_forecast=False,
            )
        self.forecast_basis = trend_model(
            F,
            num_coefficients=self.expansion_coefficient_length,
            context_length=self.context_length,
            prediction_length=self.prediction_length,
            is_forecast=True,
        )

class NBEATSNetwork(mx.gluon.HybridBlock):
    """
    The NBEATS Network as described in the paper: https://arxiv.org/abs/1905.10437.
    This does not constitute the whole NBEATS model, which is en ensemble model
    comprised of a multitude of NBEATS Networks.

    Parameters
    ----------
    prediction_length
        Length of the prediction. Also known as 'horizon'.
    context_length
        Number of time units that condition the predictions
        Also known as 'lookback period'.
    num_stacks
        The number of stacks the network should contain.
        Default and recommended value for generic mode: 30
        Recommended value for interpretable mode: 2
    num_blocks
        The number of blocks per stack.
        A list of ints of length 1 or 'num_stacks'.
        Default and recommended value for generic mode: [1]
        Recommended value for interpretable mode: [3]
    num_block_layers
        Number of fully connected layers with ReLu activation per block.
        A list of ints of length 1 or 'num_stacks'.
        Default and recommended value for generic mode: [4]
        Recommended value for interpretable mode: [4]
    widths
        Widths of the fully connected layers with ReLu activation in the blocks.
        A list of ints of length 1 or 'num_stacks'.
        Default and recommended value for generic mode: [512]
        Recommended value for interpretable mode: [256, 2048]
    sharing
        Whether the weights are shared with the other blocks per stack.
        A list of ints of length 1 or 'num_stacks'.
        Default and recommended value for generic mode: [False]
        Recommended value for interpretable mode: [True]
    expansion_coefficient_lengths
        If the type is "G" (generic), then the length of the expansion coefficient.
        If type is "T" (trend), then it corresponds to the degree of the polynomial.
        If the type is "S" (seasonal) then its not used.
        A list of ints of length 1 or 'num_stacks'.
        Default value for generic mode: [32]
        Recommended value for interpretable mode: [3]
    stack_types
        One of the following values: "G" (generic), "S" (seasonal) or "T" (trend).
        A list of strings of length 1 or 'num_stacks'.
        Default and recommended value for generic mode: ["G"]
        Recommended value for interpretable mode: ["T","S"]
    kwargs
        Arguments passed to 'HybridBlock'.
    """

    # Needs the validated decorator so that arguments types are checked and
    # the block can be serialized.
    @validated()
    def __init__(
        self,
        prediction_length: int,
        context_length: int,
        num_stacks: int,
        widths: List[int],
        num_blocks: List[int],
        num_block_layers: List[int],
        expansion_coefficient_lengths: List[int],
        sharing: List[bool],
        stack_types: List[str],
        **kwargs,
    ) -> None:
        super().__init__(**kwargs)

        self.num_stacks = num_stacks
        self.widths = widths
        self.num_blocks = num_blocks
        self.num_block_layers = num_block_layers
        self.sharing = sharing
        self.expansion_coefficient_lengths = expansion_coefficient_lengths
        self.stack_types = stack_types
        self.prediction_length = prediction_length
        self.context_length = context_length

        with self.name_scope():
            self.net_blocks: List[NBEATSBlock] = []

            # connect all the blocks correctly
            for stack_id in range(num_stacks):
                for block_id in range(num_blocks[stack_id]):
                    # in case sharing is enabled for the stack
                    params = (
                        self.net_blocks[-1].collect_params()
                        if (block_id > 0 and sharing[stack_id])
                        else None
                    )
                    # only last one does not have backcast
                    has_backcast = not (
                        stack_id == num_stacks - 1
                        and block_id == num_blocks[num_stacks - 1] - 1
                    )
                    if self.stack_types[stack_id] == "G":
                        net_block = NBEATSGenericBlock(
                            width=self.widths[stack_id],
                            num_block_layers=self.num_block_layers[stack_id],
                            expansion_coefficient_length=self.expansion_coefficient_lengths[
                                stack_id
                            ],
                            prediction_length=prediction_length,
                            context_length=context_length,
                            has_backcast=has_backcast,
                            params=params,
                        )
                    elif self.stack_types[stack_id] == "S":
                        net_block = NBEATSSeasonalBlock(
                            width=self.widths[stack_id],
                            num_block_layers=self.num_block_layers[stack_id],
                            expansion_coefficient_length=self.expansion_coefficient_lengths[
                                stack_id
                            ],
                            prediction_length=prediction_length,
                            context_length=context_length,
                            has_backcast=has_backcast,
                            params=params,
                        )
                    else:  # self.stack_types[stack_id] == "T"
                        net_block = NBEATSTrendBlock(
                            width=self.widths[stack_id],
                            num_block_layers=self.num_block_layers[stack_id],
                            expansion_coefficient_length=self.expansion_coefficient_lengths[
                                stack_id
                            ],
                            prediction_length=prediction_length,
                            context_length=context_length,
                            has_backcast=has_backcast,
                            params=params,
                        )

                    self.net_blocks.append(net_block)
                    self.register_child(
                        net_block, f"block_{stack_id}_{block_id}"
                    )

    # noinspection PyMethodOverriding,PyPep8Naming
    def hybrid_forward(self, F, past_target: Tensor, future_target: Tensor):
        if len(self.net_blocks) == 1:  # if first block is also last block
            return self.net_blocks[0](past_target)
        else:
            backcast, forecast = self.net_blocks[0](past_target)
            backcast = past_target - backcast
            # connect regular blocks (all except last)
            for i in range(1, len(self.net_blocks) - 1):
                b, f = self.net_blocks[i](backcast)
                backcast = backcast - b
                forecast = forecast + f
            # connect last block
            return forecast + self.net_blocks[-1](backcast)

    def smape_loss(self, F, forecast: Tensor, future_target: Tensor) -> Tensor:
        r"""
        .. math::

            smape = (200/H)*mean(|Y - Y_hat| / (|Y| + |Y_hat|))

        According to paper: https://arxiv.org/abs/1905.10437.
        """

        # Stop gradient required for stable learning
        denominator = F.stop_gradient(F.abs(future_target) + F.abs(forecast))
        flag = denominator == 0

        smape = (200 / self.prediction_length) * F.mean(
            (F.abs(future_target - forecast) * (1 - flag))
            / (denominator + flag),
            axis=1,
        )

        return smape

    def mape_loss(self, F, forecast: Tensor, future_target: Tensor) -> Tensor:
        r"""
        .. math::

            mape = (100/H)*mean(|Y - Y_hat| / |Y|)

        According to paper: https://arxiv.org/abs/1905.10437.
        """

        denominator = F.abs(future_target)
        flag = denominator == 0

        mape = (100 / self.prediction_length) * F.mean(
            (F.abs(future_target - forecast) * (1 - flag))
            / (denominator + flag),
            axis=1,
        )

        return mape

    def mase_loss(
        self,
        F,
        forecast: Tensor,
        future_target: Tensor,
        past_target: Tensor,
        periodicity: int,
    ) -> Tensor:
        r"""
        .. math::

            mase = (1/H)*(mean(|Y - Y_hat|) / seasonal_error)

        According to paper: https://arxiv.org/abs/1905.10437.
        """
        factor = 1 / (
            self.context_length + self.prediction_length - periodicity
        )
        whole_target = F.concat(past_target, future_target, dim=1)
        seasonal_error = factor * F.mean(
            F.abs(
                F.slice_axis(whole_target, axis=1, begin=periodicity, end=None)
                - F.slice_axis(whole_target, axis=1, begin=0, end=-periodicity)
            ),
            axis=1,
        )
        flag = seasonal_error == 0

        mase = (
            F.mean(F.abs(future_target - forecast), axis=1) * (1 - flag)
        ) / (seasonal_error + flag)

        return mase

class NBEATSTrainingNetwork(NBEATSNetwork):
    @validated()
    def __init__(self, loss_function: str, freq: str, *args, **kwargs) -> None:
        super().__init__(*args, **kwargs)
        self.loss_function = loss_function
        self.freq = freq

        # Covert frequency string like "2H" to whats called the periodicity m.
        # E.g. 12 in case of "2H" because of 12 times two hours in a day.
        self.periodicity = get_seasonality(self.freq)

        if self.loss_function == "MASE":
            assert (
                self.periodicity < self.context_length + self.prediction_length
            ), (
                "If the 'periodicity' of your data is less than 'context_length' + 'prediction_length' "
                "the seasonal_error cannot be calculated and thus 'MASE' cannot be used for optimization."
            )

    # noinspection PyMethodOverriding,PyPep8Naming
    def hybrid_forward(
        self, F, past_target: Tensor, future_target: Tensor
    ) -> Tensor:
        """

        Parameters
        ----------
        F
        past_target
            Tensor with past observations.
            Shape: (batch_size, context_length, target_dim).
        future_target
            Tensor with future observations.
            Shape: (batch_size, prediction_length, target_dim).

        Returns
        -------
        Tensor
            Loss tensor. Shape: (batch_size, ).
        """
        # future_target never used
        forecast = super().hybrid_forward(
            F, past_target=past_target, future_target=future_target
        )

        if self.loss_function == "sMAPE":
            loss = self.smape_loss(F, forecast, future_target)
        elif self.loss_function == "MAPE":
            loss = self.mape_loss(F, forecast, future_target)
        elif self.loss_function == "MASE":
            loss = self.mase_loss(
                F, forecast, future_target, past_target, self.periodicity
            )
        else:
            raise ValueError(
                f"Invalid value {self.loss_function} for argument loss_function."
            )

        return loss

class NBEATSPredictionNetwork(NBEATSNetwork):
    @validated()
    def __init__(self, *args, **kwargs) -> None:
        super().__init__(*args, **kwargs)

    # noinspection PyMethodOverriding,PyPep8Naming
    def hybrid_forward(
        self, F, past_target: Tensor, future_target: Tensor = None
    ) -> Tensor:
        """

        Parameters
        ----------
        F
        past_target
            Tensor with past observations.
            Shape: (batch_size, context_length, target_dim).
        future_target
            Not used.

        Returns
        -------
        Tensor
            Prediction sample. Shape: (batch_size, 1, prediction_length).
        """
        # future_target never used
        forecasts = super().hybrid_forward(
            F, past_target=past_target, future_target=past_target
        )

        # dimension collapsed previously because we only have one sample each:
        forecasts = F.expand_dims(forecasts, axis=1)

        # (batch_size, 1, prediction_length)
        return forecasts

max_stars_repo_name: lfywork/gluon-ts

max_issues_repo_name: lfywork/gluon-ts

max_forks_repo_name: lfywork/gluon-ts


text:
import os.path
from data.base_dataset import BaseDataset, get_transform
from data.image_folder import make_dataset
from PIL import Image
import random
import cv2 as cv
import numpy as np
import os

class ROI_transforms(object):
    def __init__(self, size=(128, 128)):
        super(ROI_transforms).__init__()
        self.size = size
        self.area = self.size[0] * self.size[1]
        self.prosepect_pmax = 0.1
        self.morph_method = [cv.MORPH_DILATE, cv.MORPH_ERODE, cv.MORPH_OPEN, cv.MORPH_CLOSE]
        self.gray_mask_method = [self.rand_mask, self.gauss_mask, self.const_mask]

    def transform_bit(self, x):
        x = self.get_shape_mask(x)
        x = self.rotate(x, angle=[0, 359])

        if random.random() > 0.5:
            x = self.scale(x, [0.5, 2], [0.5, 2])
        if random.random() > 0.5:
            x = self.random_morphological_processing(x, k_size=[3, 11])
        # x = self.align_img(x)
        x = self.crop_or_pad(x)
        _, x = cv.threshold(x, 20, 255, cv.THRESH_BINARY)
        return x

    def transform_gray(self, x):
        x = self.get_shape_mask(x)
        x = self.rotate(x, angle=[0, 359])
        if random.random() > 0.5:
            x = self.scale(x, [0.5, 2], [0.5, 2])
        # x = self.align_img(x)
        x = self.crop_or_pad(x)
        return x

    def get_new_roi(self, mask):
        """

        :param mask: (ndarray.uint8)[Height, Width]
        :return:
        """
        # 增加灰度图和二值图的判断
        # 存在20-230灰度值像素则认定为灰度图
        _, mask_dist = cv.threshold(mask, 20, 255, cv.THRESH_TOZERO)
        _, mask_dist = cv.threshold(mask_dist, 230, 255, cv.THRESH_TOZERO_INV)
        if np.count_nonzero(mask_dist) < 5:
            # 二值图处理
            # 1.mask增强
            mask = self.transform_bit(mask)
            # 2.灰度赋值
            mask = mask.astype(np.float32) / 255
            low = np.random.randint(0, 150)
            high = low + np.random.randint(50, 105)
            mask = self.gray_mask_method[np.random.randint(0, len(self.gray_mask_method) - 1)](mask, low, high)
            if np.random.random() < 0.7:
                # 平滑随机噪声
                k = np.random.randint(1, 5) * 2 + 1
                cv.GaussianBlur(mask, (k, k), k, mask)
        else:
            # 灰度图处理
            # 增强
            mask = self.transform_gray(mask)
            scale = 0.8 + 0.4 * np.random.rand()
            offset = np.random.randint(-10, 10)
            # 随机线性变换
            cv.convertScaleAbs(mask, mask, scale, offset)
        mask = mask.astype(np.uint8)
        # if mask.shape != self.size:
        #     cv.imshow("1", mask)
        #     cv.imshow("2", self.crop_or_pad(mask))
        #     cv.waitKey(0)

        return mask

    def get_shape_mask(self, x):
        if np.count_nonzero(x) < 20:
            return np.ones((np.random.randint(5, 15), np.random.randint(5, 15)), dtype=np.uint8) * 255
        Row = np.argwhere(np.sum(x, axis=0) != 0)
        Col = np.argwhere(np.sum(x, axis=1) != 0)
        x = x[np.min(Col): np.max(Col) + 1, np.min(Row): np.max(Row) + 1]
        # 控制像素数量
        while np.count_nonzero(x) > self.area * self.prosepect_pmax:
            scale = np.random.random()
            scale = scale if scale > 0.5 else 0.5
            x = cv.resize(src=x, dsize=(int(x.shape[1]*scale), int(x.shape[0]*scale)), interpolation=cv.INTER_NEAREST)
        return x

    # 旋转

    def rotate(self, x, angle=0):
        H, W = x.shape
        if isinstance(angle, list):
            assert len(angle) == 2
            angle = np.random.randint(angle[0], angle[1])

        x = np.pad(x, ((W//2, W//2), (H//2, H//2)), mode="constant", constant_values=0)
        H, W = x.shape
        m = cv.getRotationMatrix2D((W//2, H//2), angle, scale=1)
        x = cv.warpAffine(x, m, (x.shape[1], x.shape[0]))
        x = self.get_shape_mask(x)
        return x

    # 形态学处理
    def random_morphological_processing(self, x, k_size=3):
        if isinstance(k_size, list):
            k_size = np.random.randint(k_size[0], k_size[1])
        k_size = k_size // 2 * 2 + 1
        element = cv.getStructuringElement(cv.MORPH_ELLIPSE, (k_size, k_size))
        param = {"src": x, "kernel": element}
        param["op"] = self.morph_method[random.randint(0, len(self.morph_method) - 1)]
        y = cv.morphologyEx(**param)
        if np.sum(y)//255 < 10:
            return x
        return y

    # 放缩
    def scale(self, x, scaleX_factor=1, scaleY_factor=1):
        if isinstance(scaleX_factor, list):
            assert len(scaleX_factor) == 2
            scaleX_factor = scaleX_factor[0] + (scaleX_factor[1] - scaleX_factor[0]) * np.random.rand()
        if isinstance(scaleY_factor, list):
            assert len(scaleY_factor) == 2
            scaleY_factor = scaleY_factor[0] + (scaleY_factor[1] - scaleY_factor[0]) * np.random.rand()
        cv.resize(x, (int(x.shape[1] * scaleX_factor), int(x.shape[0] * scaleY_factor)), x,
                  interpolation=cv.INTER_LINEAR)
        return x

    # 回归尺寸
    # def align_img(self, x):
    #     # if np.random.random() < 0.2:
    #     #     x = self.resize(x)
    #     # else:
    #     #     x = self.crop_or_pad(x)
    #     #
    #     x = self.crop_or_pad(x)
    #     cv.threshold(x, 20, 255, cv.THRESH_BINARY, x)
    #     return x

    def resize(self, x):
        x = np.resize(x, self.size)
        return x

    def crop_or_pad(self, x):
        y = None
        cnt = 0
        while y is None or np.sum(y)//255 < 10:
            H = x.shape[0] - self.size[0]
            W = x.shape[1] - self.size[1]
            if H < 0:
                H = -H
                pad_top = random.randint(0, H)
                y = np.pad(x, ((pad_top, H - pad_top), (0, 0)), mode="constant", constant_values=0)
            else:
                crop_top = random.randint(0, H)
                y = x[crop_top: crop_top + self.size[0]]
            if W < 0:
                W = -W
                pad_left = random.randint(0, W)
                y = np.pad(y, ((0, 0), (pad_left, W - pad_left)), mode="constant", constant_values=0)
            else:
                crop_left = random.randint(0, W)
                y = y[:, crop_left: crop_left + self.size[1]]
            # crop有时只裁剪到黑色区域,此时直接resize
            if np.sum(y)//255 < 10:
                cnt += 1
                if cnt >= 5:
                    return np.resize(x, self.size).astype(np.uint8)
        return y

    # 随机mask灰度值
    def rand_mask(self, mask, low, high):
        gray_mask = np.random.randint(low, high, mask.shape) * mask
        return gray_mask

    def gauss_mask(self, mask, low, high):
        mask = self.get_shape_mask(mask)
        gauss_x = cv.getGaussianKernel(mask.shape[1], mask.shape[1])
        gauss_y = cv.getGaussianKernel(mask.shape[0], mask.shape[0])
        kyx = np.multiply(gauss_y, np.transpose(gauss_x))

        mask = mask * kyx
        Max = np.max(mask)
        Min = np.min(np.where(mask == 0, Max, mask))

        gray_mask = low + (mask - Min) / (Max - Min) * (high - low)
        gray_mask = np.where(gray_mask > 0, gray_mask, 0)
        gray_mask = self.crop_or_pad(gray_mask)
        return gray_mask

    def const_mask(self, mask, *args):
        gray_mask = mask * np.random.randint(0, 255)
        return gray_mask

# def genRandImg1(size,mask):
#     path = './gg'
#     picture_rand = os.listdir(path)
#     len_rand_picture = len(picture_rand)
#     x = random.randint(0, len_rand_picture - 1)
#     name_image = picture_rand[x]
#     picture = cv.imread(path + '/' + name_image, 0)
#     # print(type(picture))
#     picture = cv.resize(picture, (128,128))
#     # print(picture)
#     # _, mask_pict = cv.threshold(picture, 150, 255, cv.THRESH_BINARY)
#     #
#     # cv2.imshow('image',mask_pict)
#     # cv2.waitKey()
#     picture = picture.astype(np.float)
#     return picture

def get_new_image(img, gray_mask):
    gray_mask = gray_mask.astype(np.float32)
    mask = np.where(gray_mask > 0, 1, 0)
    # mask = np.where(gray_mask > 0, 255, 0)
    # mask1=mask.astype(np.uint8)
    # cv.imshow('mask',mask1)
    # cv.waitKey(5000)
    # cover
    if random.random() > 0.8:

        new_img = (img * (1 - mask) + gray_mask * mask)
    else:

        # new_img = (img * (1 - mask)) + gray_mask * mask * (255 - np.mean(img)) / 255
        new_img = (img * (1 - mask)) + mask * img * (1 + (gray_mask - 127.5) / 127.5)
    new_img = np.clip(new_img, 0, 255).astype(np.uint8)
    return new_img

def smooth_edge(new_img, mask):
    _, mask = cv.threshold(mask, 1, 255, cv.THRESH_BINARY)
    element = cv.getStructuringElement(cv.MORPH_ELLIPSE, (5, 5))
    mask_dilate = cv.morphologyEx(mask, cv.MORPH_DILATE, element)
    mask_erode = cv.morphologyEx(mask, cv.MORPH_ERODE, element)
    mask_edge = ((mask_dilate - mask_erode) / 255).astype(np.float32)
    new_img_gauss = cv.GaussianBlur(new_img, (5, 5), 5)
    return (new_img * (1 - mask_edge) + new_img_gauss * mask_edge).astype(np.uint8)

class DefectiveGenerator(object):
    def __init__(self,dir_Database,shape_Img,Limit_ROI=(20,10000),withDatabase=True):
        """

        :param dir_Database:  缺陷ROI路径
        :param shape_Img:   图片大小[height,width]
        :param Limit_ROI:   ROI外接矩形大小[lower ,upper]
        :param withDatabase:   true:从硬盘读入ROI  false：算法生成ROI
        """
        self.dir_Database=dir_Database
        self.height_Img = shape_Img[0]
        self.width_Img=shape_Img[1]
        self.lowerLimit_ROI=Limit_ROI[0]
        self.upperLimit_ROI=Limit_ROI[1]
        #
        self.roi_transform = ROI_transforms()
        #从数据库读入ROI
        self.names_ROIs,self.num_ROIs=self.loadROIs(self.dir_Database)
        if self.num_ROIs<1:
            print("the dataset is empty!")

    def loadROIs(self,dir):
        # ROIs=os.listdir(dir)
        # 递归遍历文件
        ROIs = list()
        for root, dirs, files in os.walk(dir):
            # print(root)
            for file in files:
                if file.endswith(".bmp") or file.endswith(".PNG"):
                    ROIs.append(os.path.join(root, file))

        num_ROI=len(ROIs)
        print('采用本地ROI个数为{}'.format(num_ROI))
        return ROIs,num_ROI

    def genRandImg(self,size):
        mean=random.randint(-125,125)
        fluct=random.randint(1,100)
        low=mean-fluct  #+(mean-fluct<0)*abs(mean-fluct)
        high=mean+fluct   #-(mean+fluct>255)*abs(255-(mean+fluct))
        img=np.random.randint(low,high,size)
        img=img.astype(np.float)
        return img

    def genRandImg1(self,size):
        path = './gg'
        picture_rand = os.listdir(path)
        len_rand_picture = len(picture_rand)
        x = random.randint(0, len_rand_picture - 1)
        name_image = picture_rand[x]
        picture = cv.imread(path + '/' + name_image, 0)
        # print(type(picture))
        picture = cv.resize(picture, (128,128))
        # cv.imshow('image',picture)
        # cv.waitKey()
        # print(picture)
        # _, mask_pict = cv.threshold(picture, 150, 255, cv.THRESH_BINARY)
        #
        # cv.imshow('image',picture)
        # cv.waitKey(5000)
        # picture = picture.astype(np.float)
        # cv.imshow('image',picture)
        # cv.waitKey(5000)
        return picture

    def apply(self,img):
        ROI=self.randReadROI()
        # 灰度mask处理
        # 1.最小矩形提取
        # 2.随机旋转和放缩
        # 3.尺寸回归

        # 二值mask处理
        # roi增强
        # 1.最小矩形提取形状
        # 2.随机旋转和放缩
        # 3.形态学处理
        # 4.回归尺寸
        # 返回二值掩模图

        # 返回灰度roi
        ROI_new = self.roi_transform.get_new_roi(ROI)
        ROI_new = np.where(ROI_new > 0, 1, 0).astype(np.uint8)
        #
        # cv.imshow('mask',ROI_new)
        # cv.waitKey(5000)
        randd = random.randint(0, 1)
        if randd == 0:
            img_rand = self.genRandImg([self.height_Img, self.width_Img])
        else:
            img_rand = self.genRandImg1([self.height_Img, self.width_Img])
        img_new = img.astype(np.float)
        rand = random.randint(0, 1)
        if rand == 0:
            img_new = img_new * (1 - ROI_new) + img_rand * ROI_new
        else:
            img_new = img_new + img_rand * ROI_new
        #  img_new = img_new * (1 - ROI_new) + img_rand * ROI_new
        img_new = np.clip(img_new, 0, 255).astype(np.uint8)
        # ROI_new = (ROI_new * 255).astype(np.uint8)
        # cv.imshow('mask',img_new)
        # cv.imshow('mask', img_rand)
        # cv.waitKey(5000)
        return img_new, ROI_new

        # img_new = get_new_image(img, ROI_new)
        #
        # img_new = smooth_edge(img_new, ROI_new)
        # cv.imshow("img", img)
        # cv.imshow("ROI", ROI)
        # cv.imshow("ROI_new", ROI_new)
        # cv.imshow("img_new", img_new)
        # cv.waitKey(0)

        #
        # img_rand=self.genRandImg([self.height_Img, self.width_Img])
        # img_new=img.astype(np.float)

        # rand = np.random.randint(0, 1)

        # if rand==0:
        #    img_new=img_new*(1-ROI_new)+img_rand*ROI_new
        # else:
        #     img_new = img_new + img_rand * ROI_new
      #  img_new = img_new * (1 - ROI_new) + img_rand * ROI_new
      #   img_new=np.clip(img_new, 0, 255).astype(np.uint8)
      #   ROI_new=(ROI_new*255).astype(np.uint8)
      #   return img_new, ROI_new

    def randReadROI(self):

        while(1):
            rand = random.randint(0, self.num_ROIs - 1)
            name_Img = self.names_ROIs[rand]
            img_Label = cv.imread(name_Img, 0)
            cv.threshold(img_Label, 20, 255, cv.THRESH_TOZERO, img_Label)
            if np.sum(img_Label) > 5:
                return img_Label

    def randVaryROI(self,ROI):

        return ROI

    # def randMoveROI(self,ROI):
    #     #求图像的域的大小
    #     Height_Domain =  self.height_Img
    #     Width_Domain= self.width_Img
    #     #求ROI区域的坐标
    #     Rows,Cols = np.nonzero(ROI)
    #     #求ROI区域的外接矩形大小
    #     Width_ROI=np.max(Cols)-np.min(Cols)
    #     Height_ROI=np.max(Rows)-np.min(Rows)
    #     #随机设置ROI的起始坐标
    #     Row_Upleft=random.randint(0,Height_Domain-Height_ROI-1)
    #     Col_Upleft = random.randint(0, Width_Domain - Width_ROI-1)
    #     Rows=Rows-np.min(Rows)+Row_Upleft
    #     Cols=Cols-np.min(Cols)+Col_Upleft
    #     ROI_new=np.zeros([Height_Domain,Width_Domain])
    #     ROI_new[Rows,Cols]=1
    #     return ROI_new

    # def genRandImg(self,size):
    #     mean=random.randint(-125,125)
    #     fluct=random.randint(1,100)
    #     low=mean-fluct  #+(mean-fluct<0)*abs(mean-fluct)
    #     high=mean+fluct   #-(mean+fluct>255)*abs(255-(mean+fluct))
    #     img=np.random.randint(low,high,size)
    #     img=img.astype(np.float)
    #     #
    #     return img

class RepairDataset(BaseDataset):
    """
    """

    def __init__(self, opt):
        """Initialize this dataset class.

        Parameters:
            opt (Option class) -- stores all the experiment flags; needs to be a subclass of BaseOptions
        """
        BaseDataset.__init__(self, opt)
        self.dir_A = os.path.join(opt.dataroot, opt.phase + 'A')  # create a path '/path/to/data/trainA'
        self.A_paths = sorted(make_dataset(self.dir_A, opt.max_dataset_size))   # load images from '/path/to/data/trainA'
        self.A_size = len(self.A_paths)  # get the size of dataset A
        self.B_size =  self.A_size  # get the size of dataset B
        btoA = self.opt.direction == 'BtoA'
        input_nc = self.opt.output_nc if btoA else self.opt.input_nc       # get the number of channels of input image
        output_nc = self.opt.input_nc if btoA else self.opt.output_nc      # get the number of channels of output image
        self.transform=self.get_transform()
        self.transform_A = get_transform(self.opt, grayscale=(input_nc == 1))
        self.transform_B = get_transform(self.opt, grayscale=(output_nc == 1))
        # ztb1数据换为128, 128
        self.defectGen =DefectiveGenerator("../datasets/masks", (128, 128))
        self.phase=opt.phase
    def __getitem__(self, index):
        """Return a data point and its metadata information.

        Parameters:
            index (int)      -- a random integer for data indexing

        Returns a dictionary that contains A, B, A_paths and B_paths
            A (tensor)       -- an image in the input domain
            B (tensor)       -- its corresponding image in the target domain
            A_paths (str)    -- image paths
            B_paths (str)    -- image paths
        """
        A_path = self.A_paths[index % self.A_size]  # make sure index is within then range
        if self.opt.serial_batches:   # make sure index is within then range
            index_B = index % self.B_size
        else:   # randomize the index for domain B to avoid fixed pairs.
            index_B = random.randint(0, self.B_size - 1)
       # print(A_img.size)

       #  if self.phase=="train":
       #      B_img = Image.open(A_path).convert('L')
       #      # 正负采样:
       #      if np.random.random() < 0.5:
       #          B_img = self.transform(B_img)
       #      A_img,_ = self.defectGen.apply((np.array(B_img)))
       #      A_img=Image.fromarray(A_img)
       #      # apply image transformation
       #      A = self.transform_A(A_img)
       #      B = self.transform_B(B_img)
       #      return {'A': A, 'B': B, 'A_paths': A_path}

        if self.phase == "train":
            B_img = Image.open(A_path).convert('L')
            B_img = self.transform(B_img)
            if index % 2 == 0:
                A_img, _ = self.defectGen.apply((np.array(B_img)))
                # cv.imshow('image',A_img)
                # cv.waitKey(5000)
                A_img = Image.fromarray(A_img)
                # apply image transformation
                A = self.transform_A(A_img)
                B = self.transform_B(B_img)
                # print('缺陷样本')
                return {'A': A, 'B': B, 'A_paths': A_path}
            else:
                B = self.transform_B(B_img)
                # print('正样本')
                return {'A': B, 'B': B, 'A_paths': A_path}
        elif self.phase=="test":
            A_img = Image.open(A_path).convert('L')
            A = self.transform_A(A_img)
            # B_mask_path=A_path.replace("testA","mask")
            # dirname,fname=os.path.split(B_mask_path)
            # B_mask_path=os.path.join(dirname,"groundT_"+fname)
            # B_mask= Image.open(B_mask_path).convert('L')
            # B_mask = self.transform_A(B_mask)
            # 测试级无mask,使用原图
            B_mask = A
            return {'A': A, 'B': A, 'A_paths': A_path,'B_mask': B_mask}
    def __len__(self):
        """Return the total number of images in the dataset.

        As we have two datasets with potentially different number of images,
        we take a maximum of
        """
        return self.A_size

    def get_transform(self):
        import torchvision.transforms as transforms
        l=[]
        l.append(transforms.RandomHorizontalFlip())
        l.append(transforms.RandomVerticalFlip())
        l.append(transforms.Resize([128, 128]))
       # l.append(transforms.RandomResizedCrop( 256, scale=(0.8, 1.0), ratio=(3. / 4., 4. / 3.)))
        return   transforms.Compose(l)

ext: py

reason: import numpy


text:
import struct
import numpy as np
import json
import nibabel.nicom.dicomwrappers as dcmwrapper
from pytest import approx

import autovps.dataset.siemens as siemens
from autovps.dataset import svsdata

DIM_TOL = .1

def test_orientations():
    """ Test that the axis string is correctly recovered
    """
    siemens_orientations = {'3': 'Tra', '4': 'Tra', '5': 'Tra', '6': 'Sag',
                            '7': 'Cor', '8': 'Sag', '9': 'Cor', '10': 'Tra>Sag 30.0',
                            '11': 'Tra>Cor 30.0', '12': 'Tra>Cor 30.0 >Sag 15.0',
                            '13': 'Tra>Sag 30.0 >Cor 15.0', '14': 'Tra>Sag -30.0',
                            '15': 'Tra>Cor -30.0', '16': 'Tra>Cor 30.0 >Sag -15.0',
                            '17': 'Tra>Sag 30.0 >Cor -15.0', '18': 'Tra>Sag 30.0',
                            '19': 'Tra>Cor 30.0', '20': 'Tra>Cor 30.0 >Sag 15.0',
                            '21': 'Tra>Sag 30.0 >Cor 15.0'}

    siemens_inplane_orientations = {'3': 0.0, '4': 30.0, '5': -30.0, '6': 0.0,
                                    '7': 0.0, '8': 30.0, '9': 30.0, '10': 0.0,
                                    '11': 0.0, '12': 0.0, '13': 0.0, '14': 0.0,
                                    '15': 0.0, '16': 0.0, '17': 0.0, '18': 20.0,
                                    '19': 20.0, '20': 20.0, '21': 20.0}
    for idx, ori in siemens_orientations.items():
        data = siemens.Siemens('tests/data/siemens/%s' % idx)
        data.calculate_transform()
        calc_ori = data.qform.siemens_orientation()
        # orientation string is correct
        assert ori == calc_ori[0]
        assert abs(siemens_inplane_orientations[idx] - calc_ori[1]) < DIM_TOL
        size = data.qform.get_scale()
        real_size = [20, 25, 30]
        position = data.qform.get_position()
        for i in range(3):
            assert abs(size[i]-real_size[i]) < DIM_TOL
            assert abs(position[i]) < DIM_TOL

        # in plane rotation
        assert approx(np.degrees(data.meta['VoiInPlaneRotation'])) \
            == siemens_inplane_orientations[idx]

def test_read_single():
    """Test reading a single combined DICOM
    """
    test_file = 'tests/data/siemens/3/S8457LTU_2_3_00001_00001_173218510000__1263534865.dcm'
    data_file = siemens.Siemens(test_file)
    svs_file = data_file.get_svsdata()

    data_dir = siemens.Siemens('tests/data/siemens/3')
    svs_dir = data_dir.get_svsdata()

    # is it the right size?
    assert data_file.data.shape == data_dir.data.shape == (1, 1, 1, 1, 2048)

    dcm = dcmwrapper.wrapper_from_file(test_file)
    packed = dcm.get((0x7fe1, 0x1010)).value
    data = struct.unpack("<%df" % (len(packed) / 4), packed)
    cmplx = [data[i]+data[i+1]*1j for i in range(0, len(data), 2)]

    assert (data_file.data == np.conj(cmplx)).all()
    assert (data_file.data == data_dir.data).all()

def test_read_uncombined():
    """Test reading a directory of dicoms
    """
    data = siemens.Siemens('tests/data/siemens/eja_svs_press_uncombined')
    data.get_svsdata()

    assert data.data.shape == (52, 32, 1, 1, 2048)
    for i in range(data.data.shape[0]):
        for j in range(data.data.shape[1]):
            assert data.data[i, j, 0, 0, -1] != 0

def test_read_unaveraged():
    """Test reading a directory of dicoms
    """
    data = siemens.Siemens('tests/data/siemens/eja_svs_press_combined_noave')
    data.get_svsdata()

    assert data.data.shape == (1, 8, 1, 1, 2048)
    for i in range(data.data.shape[0]):
        for j in range(data.data.shape[1]):
            assert data.data[i, j, 0, 0, -1] != 0

def test_sidecar():
    """Test creating a BIDS sidecar
    """
    # TODO: Check value consistency
    data = siemens.Siemens('tests/data/siemens/eja_svs_press_uncombined')
    with open('test.json', 'w+') as fp:
        json.dump(data.get_sidecar(), fp, indent=3)

    with open('test.json') as fp:
        bids = json.load(fp)

    # save the transform
    np.savetxt('sub-test_from-Device_to-orig_mode-image_xfm.mat', data.qform.get_matrix())

max_stars_repo_name: rhancockn/autovps

max_issues_repo_name: rhancockn/autovps

max_forks_repo_name: rhancockn/autovps


text:
#!/usr/bin/python
# -*- coding: utf-8 -*-
import numpy as np
from scipy.optimize import brentq
from scipy.interpolate import interp1d

def get_critical_point(eosmat, x, y, x_min=None, x_max=None, y_min=None, y_max=None, extra_fields={}):
    from scipy.spatial.distance import cdist
    # get a contours for Pt_DT.dFx == 0 and Pt_DT.dFxx == 0
    Pt = get_contour(eosmat.Pt_DT.dFx, x, y, 0.0, x_min=x_min, x_max=x_max, y_min=y_min, y_max=y_max)
    Pt_x = get_contour(eosmat.Pt_DT.dFx, x, y, 0.0, x_min=x_min, x_max=x_max, y_min=y_min, y_max=y_max, extra_fields=[eosmat.Pt_DT])
    Pt_xx = get_contour(eosmat.Pt_DT.dFxx, x, y, 0.0, x_min=x_min, x_max=x_max, y_min=y_min, y_max=y_max, extra_fields=[eosmat.Pt_DT])

    # out of all the results select the point with the highest temperature

    rho_c, temp_c = get_curve_intersection(Pt_x[:,0], Pt_x[:,1], Pt_xx[:,0], Pt_xx[:,1], log=True)
    idx = np.argmax(temp_c)

    rho_c = rho_c[idx]
    temp_c = temp_c[idx]

    pres_c = float(eosmat.Pt_DT(rho_c, temp_c))
    #if pres_c < 0:
    #    raise ValueError("The critical point can't have a negative pressure!")

    out = {'pres_c': pres_c, 'rho_c': rho_c, 'temp_c': temp_c}
    for key, func in extra_fields.iteritems():
        out[key] = float(func(out['rho_c'], out['temp_c']))
    return out

def saturation_curve(tab, rho, temp, extra_fields={}):
    crit_pt = get_critical_point(tab, rho, temp)
    rho_c = crit_pt['rho_c']
    temp_c = crit_pt['temp_c']
    temp = temp[temp<=temp_c]
    sat_l = []
    sat_v = []
    for idx, temp_val in enumerate(temp):
       res  = maxwell_construction(tab, rho, temp_val, rho_c, full_output=False)
       if res:
           sat_l.append([res['rho_l'], temp_val])
           sat_v.append([res['rho_v'], temp_val])
    sat_l.append([rho_c, temp_c])
    sat_v.append([rho_c, temp_c])

    sat_l = np.array(sat_l)
    sat_v = np.array(sat_v)

    Pneg = get_contour(tab.Pt_DT, rho, temp)
    mask =  (Pneg[:,1]<sat_l[:,1].min())*(Pneg[:,0]>rho_c)

    sat_l = np.concatenate((Pneg[mask,:], sat_l), axis=0)

    sat_l = {'rho': sat_l[:,0], 'temp': sat_l[:,1]}
    sat_v = {'rho': sat_v[:,0], 'temp': sat_v[:,1]}

    out = {'sat_l': sat_l, 'sat_v': sat_v}
    for maj_key in out:
        for key, func in extra_fields.iteritems():
            out[maj_key][key] = func(out[maj_key]['rho'], out[maj_key]['temp'])

    return  out

def maxwell_construction(tab, rho, temp_val, rho_c, full_output=False):
    """
    Compute Maxwell constractions on the isotherm
    Parameters:
    -----------
      - rho [ndarray] should be increasing
      - pres [ndarray]
      - temp [float]
    """
    if (np.diff(rho)<0).any():
        raise ValueError('Density array should be increasing!')
    temp = np.ones(rho.shape)*temp_val
    mask = rho>rho_c
    pres = tab.Pt_DT(rho, temp)
    A = tab.At_DT(rho, temp)
    G = A + pres/rho

    pres_m_eq, G_m_eq = get_curve_intersection(pres[mask], G[mask], pres[~mask], G[~mask])
    if len(pres_m_eq) != 1:
        return None

    dp_func = lambda x:  float(tab.Pt_DT(x, temp_val)) - pres_m_eq
    roots = global_roots(dp_func, rho)
    if len(roots) != 3:
        # must have exactly 3 roots!
        return None
    rho_l, rho_r = roots[[0,-1]]
    out = {"pres_m_eq": pres_m_eq, 'rho_v': rho_l, 'rho_l': rho_r}
    if full_output:
        pres_m = np.array(pres)
        eint_m = tab.Ut_DT(rho, temp)
        eint_m_l = float(tab.Ut_DT(rho_l, temp_val))
        eint_m_r = float(tab.Ut_DT(rho_r, temp_val))
        mask = (rho>=rho_l)*(rho<=rho_r)
        pres_m[mask] = pres_m_eq
        eint_m[mask] = (eint_m_l + (eint_m_r - eint_m_l)*(1./rho - 1/rho_l)/(1./rho_r - 1./rho_l))[mask]
        out.update(**{'pres_m': pres_m, 'eint_m': eint_m})

    return out

def global_roots(func, x, extra_args=None):
    """
    Parameters:
    -----------
      - func: function of one argument
      - x [ndarray]: points at which the function will be evaluated
    """
    roots = []
    for idx in range(len(x)-1):
        x_l, x_r = x[idx], x[idx+1]
        if func(x_l)*func(x_r)<0:
                roots.append(brentq(func, x_l, x_r))
    return np.array(roots)

def get_curve_intersection(x1, y1, x2, y2, log=False):
    """Find all intersection of curves 1 and 2 defined by points x, y
    The curve 1 is assumed to be such that for every x value there is only one y value.
    """
    if log:
        x1 = np.log10(x1)
        y1 = np.log10(y1)
        x2 = np.log10(x2)
        y2 = np.log10(y2)

    itp = interp1d(x1, y1, kind='linear',  bounds_error=False)
    x_root = []
    y_root = []
    for idx in range(len(y2)-1):
        x_l, x_r = x2[[idx, idx+1]]
        y2_l, y2_r = y2[[idx, idx+1]]
        a = (y2_r - y2_l)/(x_r-x_l)
        b = y2_l - a*x_l
        err = lambda x: itp(x) - (a*x+b)
        if err(x_l)*err(x_r)<0: # there is an intersection
            x_sol = brentq(err, x_l, x_r)
            y_sol = a*x_sol+b
            x_root.append(x_sol)
            y_root.append(y_sol)
        else:
            continue
    x_root = np.array(x_root)
    y_root = np.array(y_root)
    if log:
        x_root = np.power(10, x_root)
        y_root = np.power(10, y_root)
    return x_root, y_root

def get_contour(func, x, y, level=0.0, x_min=None, x_max=None, y_min=None, y_max=None, extra_fields=[]):
    """
    Calculate contour for a  variable

    Parameters:
    -----------
      - func : for instace EosMaterial.Pt_DT or EosMaterial.Pt_DT.dFxx
      - x [ndarray] (N,) e.g. density 1D array
      - y [ndarray] (N,) e.g. temperature  1D array
    Returns:
    --------
       ndarray (N, 2+len(extra_fields)) the columns correspond to x, y, and extra_fields
    """
    from skimage.measure import find_contours
    x_mask = ~np.isnan(x)
    y_mask = ~np.isnan(y)
    if x_min:
        x_mask *= (x>x_min)
    if x_max:
        x_mask *= (x<x_max)
    if y_min:
        y_mask *= (y>y_min)
    if y_max:
        y_mask *= (y<y_max)

    x = x[x_mask]
    y = y[y_mask]
    X, Y = np.meshgrid(x, y, indexing='xy')

    M =  func(X, Y).T
    res = find_contours(M, level=level, fully_connected='low')
    out = []
    for el in res:
        x_contour = el[:,0]
        y_contour = el[:,1]
        F = [ np.interp(x_contour, range(len(x)), x),
              np.interp(y_contour, range(len(y)), y)]
        for field in extra_fields:
            F.append(field(x_contour, y_contour))
        out.append(np.array(F).T)

    return sorted(out, key=lambda x: x.shape[0])[-1]

max_stars_repo_name: rth/pyeospac

max_issues_repo_name: rth/pyeospac

max_forks_repo_name: luli/pyeospac

reason: import numpy,from scipy