text:
#!/usr/bin/env python
"""Complex domain coloring."""

import numpy as np
import matplotlib
matplotlib.use('Agg')
from matplotlib import pyplot as plt
from matplotlib.colors import hsv_to_rgb
from tqdm import tqdm
from celluloid import Camera

fig = plt.figure()
# hack to remove border
# https://stackoverflow.com/a/37810568/744520
fig.set_size_inches(4, 4, forward=False)
ax = plt.Axes(fig, [0, 0, 1, 1])
ax.set_axis_off()
fig.add_axes(ax)
camera = Camera(fig)

for a in tqdm(np.linspace(0, 2 * np.pi, 30, endpoint=False)):
    x = np.linspace(-3, 3, 800)
    X, Y = np.meshgrid(x, x)
    x = X + 1j * Y
    y = (x ** 2 - 2.5) * (x - 2.5 * 1j) * (x + 2.5 * 1j) * (x - 2 - 1j) ** 2 / ((x - np.exp(1j * a)) ** 2 * (x - np.exp(1j * 2 * a)) ** 2)

    H = np.angle(y) / (2 * np.pi) + .5
    r = np.log2(1. + np.abs(y))
    S = (1. + np.abs(np.sin(2. * np.pi * r))) / 2.
    V = (1. + np.abs(np.cos(2. * np.pi * r))) / 2.

    rgb = hsv_to_rgb(np.dstack((H, S, V)))
    ax.imshow(rgb)
    camera.snap()
animation = camera.animate(interval=50, blit=True)
animation.save(
    'complex.mp4',
    dpi=100,
    savefig_kwargs={
        'frameon': False,
        'pad_inches': 0
    }
)

max_issues_repo_name: StevE-Ong/celluloid

max_forks_repo_name: StevE-Ong/celluloid


text:
from typing import (
    Callable,
    List,
    NamedTuple,
    Optional,
    Tuple,
    Union,
)

import numpy as np

Table = np.ndarray  # 2D matrix
Data = Table
Centroids = Table
IntLabels = np.ndarray
SegmentationMethod = Callable[[Data], Tuple[IntLabels, Centroids]]
Clustering = Union["divik.cluster.GAPSearch", "divik.cluster.DunnSearch"]

class DivikResult(NamedTuple):
    """Result of DiviK clustering"""

    clustering: Clustering
    """Fitted automated clustering estimator"""
    feature_selector: "divik.feature_selection.StatSelectorMixin"
    """Fitted feature selector"""
    merged: IntLabels
    """Recursively merged clustering labels"""
    subregions: List[Optional["DivikResult"]]
    """DivikResults for all obtained subregions"""

ext: py

max_stars_repo_name: Hirni-Meshram/divik

max_issues_repo_name: Hirni-Meshram/divik

max_forks_repo_name: Hirni-Meshram/divik

reason: import numpy


text:
#!/usr/bin/env python
"""Run an identical twin flux inversion OSSE with real data.

Use xarray/dask to grab influence functions and priors from netCDF
files.
"""
from __future__ import print_function, division, unicode_literals

import itertools
import datetime
import textwrap
import os.path
import glob
import sys

import dask.array as da
import pandas as pd
import dateutil.tz
import numpy as np
import cf_units
import netCDF4
import sparse
import xarray
import wrf

try:
    THIS_DIR = os.path.dirname(__file__)
except NameError:
    THIS_DIR = os.getcwd()

sys.path.insert(0, os.path.join(
    THIS_DIR, "..", "src"))
sys.path.append(THIS_DIR)

import atmos_flux_inversion.optimal_interpolation
import atmos_flux_inversion.variational
import atmos_flux_inversion.correlations
import atmos_flux_inversion.covariances
from atmos_flux_inversion.util import kronecker_product
from atmos_flux_inversion.linalg import asarray, kron
from atmos_flux_inversion.noise import gaussian_noise
import cf_acdd

INFLUENCE_PATHS = ["/mc1s2/s4/dfw5129/data/LPDM_2010_fpbounds/"
                   "ACT-America_trial5/2010/01/GROUP1",
                   "/mc1s2/s4/dfw5129/data/LPDM_2010_fpbounds/"
                   "candidacy_more_towers/2010/01/GROUP1"]
PRIOR_PATH = "/mc1s2/s4/dfw5129/inversion_code/data_files"
OBS_PATH = "/mc1s2/s4/dfw5129/inversion"

FLUX_INTERVAL = 6
"""The interval at which fluxes become available in hours.

Fluxes are usually integrated forward from hourly input, but I can't
solve for that in a reasonable timeframe.

This determines both the input and the output time resolution as well
as what the inversion solves for.

Note
----
Must divide twenty-four.
"""
FLUX_RESOLUTION = 27
"""Resolution of fluxes and influence functions in kilometers."""
UNCERTAINTY_RESOLUTION_REDUCTION_FACTOR = 16
"""How much coarser uncertainty is than mean estimate in the x direction.

If we compute the uncertainty at full resolution, the resulting file
is huge, starting in the hundreds of terabytes for a month.
Coarsening the resolution for the uncertainties allows us to still
report uncertainties within current computing constraints.
"""
# 4: 2h49 wall 132GiB mem 3h57 cpu
# 3: 5h11 wall 140GiB mem 6h27 cpu
UNCERTAINTY_FLUX_RESOLUTION = (UNCERTAINTY_RESOLUTION_REDUCTION_FACTOR *
                               FLUX_RESOLUTION * 1e3)
"""Resolution of posterior uncertainties in meters."""
UNCERTAINTY_TEMPORAL_RESOLUTION = "7D"
"""The resolution at which the uncertainty is calculated and saved.

Higher resolution means the uncertainties will be more accurate.
"""

# Linear interpolation in space
OBS_FILES = glob.glob(os.path.join(
    OBS_PATH,
    "2010_07_[45]tower_{inter:02d}hr_{res:03d}km_"
    "LPDM_concentrations?.nc".format(
        inter=FLUX_INTERVAL, res=FLUX_RESOLUTION)))
CORR_FUN = "exp"
CORR_LEN = 200
TIME_CORR_FUN = "exp"
TIME_CORR_LEN = 21
FLUX_FILES = glob.glob(os.path.join(
    PRIOR_PATH,
    ("2010-07_osse_bio_priors_{interval:1d}h_{res:02d}km_noise_"
     "{corr_fun:s}{corr_len:d}km_"
     "{corr_fun_time:s}{corr_len_time:d}d_exp3h.nc").format(
        interval=FLUX_INTERVAL, corr_fun=CORR_FUN, corr_len=CORR_LEN,
        corr_fun_time=TIME_CORR_FUN, corr_len_time=TIME_CORR_LEN,
        res=FLUX_RESOLUTION)))
FLUX_FILES.sort()
OBS_FILES.sort()
INFLUENCE_FILES = [
    name
    for path in INFLUENCE_PATHS
    for name in glob.glob(os.path.join(
        path,
        "LPDM_2010_01_{flux_interval:02d}hrly_{res:03d}km_molar_footprints.nc4"
        .format(flux_interval=FLUX_INTERVAL, res=FLUX_RESOLUTION)))]

print("Flux files", FLUX_FILES)
print("Influence Files", INFLUENCE_FILES)

# tracers are:
#  diurnal bio, fossil, ocean, biomass burn, biofuel,
#  ship, posterior bio, empty, prior bio, empty
#    CMS posterior mole fractions are added to the first empty tracer,
#    and note the modified CMS posterior mole fraction file used for the
#     last tracer (prefix 'tsq')
TRUE_FLUX_NAME = "E_TRA7"
TRACER_NAME = "tracer_7_LPDM"
PRIOR_FLUX_NAME = TRUE_FLUX_NAME + "_noisy"
HOURS_PER_DAY = 24
DAYS_PER_WEEK = 7
FLUX_WINDOW = HOURS_PER_DAY * DAYS_PER_WEEK * 2
"""How long fluxes considered to have an influence.

Measured in hours.
Implemented by slicing out this much of the stored influence functions
for use as the linearized observation operator in the inversion.
"""
OBS_HOURS = (datetime.time(12), datetime.time(16))
"""Which observation times will be used in the inversion.

Assumed to be local solar. Should give afternoon hours for the domain.

I really hope I can assume this doesn't depend on latitude. That would
make this much more complicated.
"""
OBS_TIMES_PER_DAY = OBS_HOURS[1].hour - OBS_HOURS[0].hour
"""Observations used per site per day."""
CO2_MOLAR_MASS = 16 * 2 + 12.01
"""Molar mass of CO2 (g/mol).

Used to convert WRF fluxes to units expected by observation operator.
"""
DAYS_DROPPED_FROM_END = 1
"""Currently 1 to avoid problems with lack of fluxes in August."""
OBS_DAYS = 30
N_REALIZATIONS = 80
#  1    9m55 (80 realizations)    9m46
# 30 2h48m40 (80 realizations) 3h56m41
"""Number of days of obs to use."""
OBS_WINDOW = OBS_DAYS * OBS_TIMES_PER_DAY
"""Number of observation times."""
CO2_MOLAR_MASS_UNITS = cf_units.Unit("g/mol")
FLUX_UNITS = cf_units.Unit("mol/m^2/s")
BAD_SITES = ("WGC", "OSI")

FLUX_CHUNKS = HOURS_PER_DAY * 8 // FLUX_INTERVAL
# 48 51s
#  4 54s
#  2 4m23
"""How many flux times to treat at once.

Must be a multiple of day length.
"""
OBS_CHUNKS_ALL = 24
"""How many observations to load at once.

Must allow a few chunks of the influence function to be in memory at
once.  Will need to extend this if (N_OBS_TIMES * N_SITES) ** 2 arrays
stop fitting nicely in memory.
"""
OBS_CHUNKS_USED = 96
"""Obs chunks treated at once in inversion code.

Must also allow a few chunks to be loaded at once.
"""
REALIZATION_CHUNK = 20
NC_ENGINE = "netcdf4"

############################################################
# Utility functions.
def grouper(iterable, n, fillvalue=None):
    """Collect data into fixed-length chunks or blocks."""
    # grouper('ABCDEFG', 3, 'x') --> ABC DEF Gxx"
    args = [iter(iterable)] * n
    return itertools.zip_longest(*args, fillvalue=fillvalue)

def sort_key_to_consecutive(sequence):
    """Turn a list of sort keys into a list of consecutive numbers.

    Parameters
    ----------
    sequence: collections.Sequence
        The list to sort.

    Returns
    -------
    tuple
        ``argsort(sequence)``, I think
    """
    items = list(enumerate(sorted(sequence)))

    items.sort(key=lambda item: sequence.index(item[1]))
    return tuple(item[0] for item in items)

def flush_output_streams():
    """Flush stdout and stderr."""
    sys.stdout.flush()
    sys.stderr.flush()

def write_progress_message(msg):
    """Write the message to stdout with time.

    Parameters
    ----------
    msg: str
    """
    flush_output_streams()
    print(datetime.datetime.now(UTC).strftime("%c"), msg)
    flush_output_streams()

UTC = dateutil.tz.tzutc()
SECONDS_PER_HOUR = 3600

############################################################
# Read sizes from influence function file
TEST_DS = netCDF4.Dataset(INFLUENCE_FILES[0])

NX = len(TEST_DS.dimensions["dim_x"])
NY = len(TEST_DS.dimensions["dim_y"])
N_TIMES_BACK = len(TEST_DS.dimensions["time_before_observation"])

N_SITES = len(TEST_DS.dimensions["site"])
N_OBS_TIMES = len(TEST_DS.dimensions["observation_time"])

TEST_DS.close()
del TEST_DS

if N_TIMES_BACK < FLUX_WINDOW / FLUX_INTERVAL:
    raise ValueError("FLUX_WINDOW too long for file")

N_GRID_POINTS = NY * NX
STATE_SIZE = N_GRID_POINTS * FLUX_WINDOW

OBS_VEC_SIZE = N_SITES * OBS_WINDOW
OBS_VEC_TOTAL_SIZE = N_SITES * N_OBS_TIMES

############################################################
# Read influence functions
print("Days of obs used:", OBS_DAYS)

# Quadratic form in save_sum
# OPTIMAL_ELEMENTS 1e4, not changed in save_sum
# obs_chunk: 24: 96 flux_chunk 3 days
# 2 week flux window, 40 realizations in loaded noise
# 1* OPTIMAL_ELEMENTS**2
# OPTIMAL_ELEMENTS 5e4
# flux_chunk 3 days, obs_chunk 24: 96
# run under debugger
#  2 days obs:  2.5 min, 2.3min
#  4 days obs:  4.3 min
#  8 days obs:  9.5 min
# 16 days obs: 24.5 min
# 31 days obs: 83.3 min
#   8m for Kd, 75 to find HBH^T
INFLUENCE_DATASET = xarray.open_mfdataset(
    INFLUENCE_FILES,
    chunks=dict(observation_time=OBS_CHUNKS_ALL, site=1,
                time_before_observation=FLUX_CHUNKS,
                dim_y=NY, dim_x=NX),
    engine=NC_ENGINE,
).isel(
    observation_time=slice(DAYS_DROPPED_FROM_END * HOURS_PER_DAY,
                           (OBS_DAYS + DAYS_DROPPED_FROM_END) * HOURS_PER_DAY),
    time_before_observation=slice(0, FLUX_WINDOW // FLUX_INTERVAL))
OBS_TIME_INDEX = (INFLUENCE_DATASET.indexes["observation_time"].round("S") +
                  datetime.timedelta(days=181))
TIME_BACK_INDEX = (
    INFLUENCE_DATASET.indexes["time_before_observation"].round("S"))
INFLUENCE_DATASET.coords["observation_time"] = OBS_TIME_INDEX
INFLUENCE_DATASET.coords["time_before_observation"] = TIME_BACK_INDEX

FLUX_TIMES = (INFLUENCE_DATASET.coords["flux_time"] +
              np.array(datetime.timedelta(days=181),
                       dtype='m8[ns]'))
INFLUENCE_DATASET.coords["flux_time"] = FLUX_TIMES

INFLUENCE_FUNCTIONS = INFLUENCE_DATASET.H.astype(np.float32)
# Use site names as index/dim coord for site dim
INFLUENCE_FUNCTIONS.coords["site"] = np.char.decode(
    INFLUENCE_FUNCTIONS["site_names"].values, "ascii")

INFLUENCE_FUNCTIONS.coords["observation_time"] = OBS_TIME_INDEX

# NB: Remember to change frequency and time zone as necessary.
FLUX_START = (OBS_TIME_INDEX[-1] - TIME_BACK_INDEX[-1]).replace(hour=0)
if OBS_TIME_INDEX[0].hour != 0:
    FLUX_END = OBS_TIME_INDEX[0].replace(hour=0) + datetime.timedelta(days=1)
else:
    FLUX_END = OBS_TIME_INDEX[0]
FLUX_TIMES_INDEX = pd.date_range(
    FLUX_START, FLUX_END,
    freq="{flux_interval:d}H".format(flux_interval=FLUX_INTERVAL),
    tz="UTC", closed="right",
    name="flux_times")
N_FLUX_TIMES = len(FLUX_TIMES_INDEX)
assert N_FLUX_TIMES % 4 == 0

############################################################
# Set some constants based on the WRF file
with netCDF4.Dataset(FLUX_FILES[0]) as ds:
    WRF_PROJECTION = wrf.util.getproj(**wrf.util.get_proj_params(ds))
    WRF_CRS = WRF_PROJECTION.cartopy()

    # Alternate method
    # Currently broken for some reason.
    # WRF_CRS = wrf.get_cartopy(ds)

WRF_TOWER_COORDS = WRF_CRS.transform_points(
    WRF_CRS.as_geodetic(),
    INFLUENCE_FUNCTIONS.coords["site_lons"].values,
    INFLUENCE_FUNCTIONS.coords["site_lats"].values,
    INFLUENCE_FUNCTIONS.coords["site_heights"].values)
OBS_ROUGH_LEVEL = 11
"""Because I don't feel like destaggering the geopotential."""
OBS_ROUGH_SIGMA = 0.9486
"""Also eyeballed from a single time."""

print(datetime.datetime.now(UTC).strftime("%c"),
      "Have constants, getting priors")
############################################################
# Read prior fluxes
FLUX_DATASET = xarray.open_mfdataset(
    FLUX_FILES,
    chunks=dict(dim_x=NX, dim_y=NY,
                flux_time=FLUX_CHUNKS,
                realization=REALIZATION_CHUNK),
    concat_dim="flux_time",
    engine=NC_ENGINE,
).isel(realization=slice(0, N_REALIZATIONS))

OBS_FILES_4 = [name for name in OBS_FILES if "4tower" in name]
OBS_FILES_5 = [name for name in OBS_FILES if "5tower" in name]

OBS_DATASET_4 = xarray.open_mfdataset(
    OBS_FILES_4,
    chunks=dict(forecast_reference_time=OBS_CHUNKS_USED),
    engine=NC_ENGINE,
)
OBS_DATASET_5 = xarray.open_mfdataset(
    OBS_FILES_5,
    chunks=dict(forecast_reference_time=OBS_CHUNKS_USED),
    engine=NC_ENGINE,
)
OBS_DATASET = xarray.concat([OBS_DATASET_4, OBS_DATASET_5],
                            dim="dim1")

write_progress_message("Have obs, normalizing")

wrf_times = OBS_DATASET.indexes["forecast_reference_time"].round("S")
OBS_DATASET.coords["forecast_reference_time"] = wrf_times

print(OBS_DATASET.dims, OBS_DATASET.coords)
OBS_DATASET.coords["site"] = list(
    map(lambda x: x.decode("ascii"),
        OBS_DATASET["name_of_observation_site"].values))
OBS_DATASET = (
    OBS_DATASET
    .set_index(dim1="site")
    .rename(dict(dim1="site"))
)
del OBS_DATASET.coords["name_of_observation_site"]
print(OBS_DATASET.dims, OBS_DATASET.coords)
# Assign a few more coords and pull out only the fluxes we need.
FLUX_DATASET = FLUX_DATASET.sel(flux_time=FLUX_TIMES_INDEX.tz_convert(None))
N_REALIZATIONS = len(FLUX_DATASET.indexes["realization"])

WRF_DX = FLUX_DATASET.attrs["DX"]

TRUE_FLUXES = FLUX_DATASET.get(["E_TRA{:d}".format(i + 1)
                                for i in range(10)])
TRUE_FLUXES_MATCHED = TRUE_FLUXES
# for flux_part, flux_orig in zip(
#         TRUE_FLUXES_MATCHED.data_vars.values(),
#         TRUE_FLUXES.data_vars.values()):
#     unit = (cf_units.Unit(flux_orig.attrs["units"]) *
#             CO2_MOLAR_MASS_UNITS)
#     # For whatever reason this is backwards from the conversion
#     # factors used elsewhere.
#     flux_part *= (unit / FLUX_UNITS).convert(1, 1)
#     flux_part.attrs["units"] = str(FLUX_UNITS)

PRIOR_FLUXES = FLUX_DATASET.get(["E_TRA{:d}_noisy".format(i + 1)
                                 for i in (6,)])
PRIOR_FLUXES_MATCHED = PRIOR_FLUXES
# for flux_part, flux_orig in zip(
#         PRIOR_FLUXES_MATCHED.data_vars.values(),
#         PRIOR_FLUXES.data_vars.values()):
#     unit = (cf_units.Unit(flux_orig.attrs["units"]) *
#             CO2_MOLAR_MASS_UNITS)
#     # For whatever reason this is backwards from the conversion
#     # factors used elsewhere.
#     flux_part *= (unit / FLUX_UNITS).convert(1, 1)
#     flux_part.attrs["units"] = str(FLUX_UNITS)

WRF_OBS = OBS_DATASET.get(
    ["tracer_{:d}_LPDM".format(i + 1)
     for i in (1, 6, 8)])
WRF_OBS_MATCHED = WRF_OBS.rename(dict(
    forecast_reference_time="observation_time"))
WRF_OBS_SITE = WRF_OBS_MATCHED
print(WRF_OBS_MATCHED.dims, WRF_OBS_MATCHED.coords)

WRF_OBS_START = WRF_OBS_MATCHED.indexes["observation_time"][0]
WRF_OBS_INTERVAL = (WRF_OBS_START -
                    WRF_OBS_MATCHED.indexes["observation_time"][1])

write_progress_message("Getting solar times")
############################################################
# Get time zone representing local solar time for each site
LOCAL_TIME_ZONES = list(map(
    dateutil.tz.tzoffset,
    ("Local time for {name!s}".format(name=name)
     for name in INFLUENCE_FUNCTIONS.coords["site"].values),
    [round(offset_hr * SECONDS_PER_HOUR)
     for offset_hr in INFLUENCE_FUNCTIONS.coords["site_lons"].values / 15]))

print(datetime.datetime.now(UTC).strftime("%c"),
      "Converting influence function to have alignment necessary for H")
############################################################
# Do the inversion
############################################################
# The Kalman filter version subtracts an hour from these. Why?
obs_times = (INFLUENCE_FUNCTIONS.indexes["observation_time"][::-1])

############################################################
# Align observation flux times
# Take care of missing obs
# Also subsetting for late afternoon steady convective boundary layer
# takes roughly 11 min.
# Takes 6 min with everything in memory
# Two minutes with big chunks?
site_obs_index = []
write_progress_message("Selecting observations")
for i, site in enumerate(INFLUENCE_FUNCTIONS.indexes["site"]):
    if site in BAD_SITES:
        continue
    local_times = pd.Index([
        obs_time.tz_localize(UTC)
        .tz_convert(LOCAL_TIME_ZONES[i])
        for obs_time in obs_times], name="local_times")

    # np.where is designed for multi-dimensional arrays
    site_obs_index.extend(zip(
        np.where((OBS_HOURS[0] < local_times.time) &
                 (local_times.time < OBS_HOURS[1]))[0],
        itertools.repeat(site)))
obs_index, site_index = zip(*site_obs_index)
site_index = np.array(list(site_index))
pd_obs_index = obs_times[list(obs_index)]
site_obs_pd_index = pd.MultiIndex.from_tuples(
    list(zip(site_index, pd_obs_index)), names=("site", "observation_time"))

print(datetime.datetime.now(UTC).strftime("%c"),
      "Aligning flux times in influence function")
flush_output_streams()
dimension_order = [item for item in INFLUENCE_FUNCTIONS.dims]
dimension_order.insert(dimension_order.index("time_before_observation"),
                       "flux_time")
dimension_order.insert(dimension_order.index("observation_time"),
                       "observation")
dimension_order = tuple(dimension_order)
aligned_influences = xarray.concat(
    [here_infl.set_index(
        time_before_observation="flux_time").rename(
            dict(time_before_observation="flux_time"))
     for here_infl in INFLUENCE_FUNCTIONS.sel(
         site=xarray.DataArray(site_index, name="observation",
                               dims=dict(observation=site_obs_pd_index)),
         observation_time=xarray.DataArray(
             pd_obs_index, name="observation",
             dims=dict(observation=site_obs_pd_index)
         )
    )],
    "observation").set_index(
    observation=("observation_time", "site"))
print(datetime.datetime.now(UTC).strftime("%c"),
      "Aligned flux times in influence function, "
      "aligning fluxes with influence function")
flush_output_streams()
aligned_influences, aligned_true_fluxes, aligned_prior_fluxes = (
    xarray.align(
        aligned_influences,
        TRUE_FLUXES_MATCHED[TRUE_FLUX_NAME],
        PRIOR_FLUXES_MATCHED[PRIOR_FLUX_NAME],
        exclude=("dim_x", "dim_y"),
        join="outer", copy=False))
print(datetime.datetime.now(UTC).strftime("%c"),
      "Aligned fluxes and influence function")
flush_output_streams()
aligned_true_fluxes = aligned_true_fluxes.transpose(
    "flux_time", "dim_y", "dim_x")
aligned_prior_fluxes = aligned_prior_fluxes.transpose(
    "flux_time", "dim_y", "dim_x", "realization")
aligned_influences = aligned_influences.transpose(
    "observation", "flux_time", "dim_y", "dim_x")
write_progress_message("Rechunked to square")
aligned_influences = aligned_influences.fillna(0)
aligned_true_fluxes.astype(np.float32).load()
aligned_prior_fluxes.astype(np.float32).load()
aligned_influences.astype(np.float32).load()
write_progress_message("Loaded data")

# 23 min for seven towers over a month
# 11 min next run
# This includes the realignment time deferred with dask above
sparse_influences = sparse.COO(aligned_influences.values)
aligned_influences.data = da.asarray(sparse_influences)

print(datetime.datetime.now(UTC).strftime("%c"), "Converted to COO")
flush_output_streams()

posterior_var_atts = aligned_prior_fluxes.attrs.copy()
posterior_var_atts.update(dict(
    long_name="posterior_fluxes",
    units=PRIOR_FLUXES_MATCHED[PRIOR_FLUX_NAME].attrs["units"],
    description="posterior fluxes using dask for a month",
    origin="OI using dask for a month",
    prior_flux_name=PRIOR_FLUX_NAME,
    flux_window=FLUX_WINDOW,
    observation_window=OBS_WINDOW,
    ancillary_variables="reduced_posterior_covariance",
))
increment_var_atts = aligned_prior_fluxes.attrs.copy()
increment_var_atts.update(dict(
    long_name="flux_increment",
    units=PRIOR_FLUXES_MATCHED[PRIOR_FLUX_NAME].attrs["units"],
    description="Change from prior to posterior using dask",
    origin="OI and dask",
    flux_window=FLUX_WINDOW,
    observation_window=OBS_WINDOW))
posterior_global_atts = cf_acdd.global_attributes_dict()
posterior_global_atts.update(dict(
    title="Posterior fluxes",
    summary="Posterior fluxes",
    creator_institution="PSU Department of Meteorology",
    product_version="v0.0.0.dev0",
    cdm_data_type="grid",
    institution="PSU Department of Meteorology",
    source="Test inversion using OI for a monthlong window",
    external_variables="reduced_posterior_covariance reduced_prior_covariance",
))

############################################################
# Define correlation constants and get covariances
# Constants for inversion
write_progress_message("Getting covariances")
CORRELATION_LENGTH = 1000
GRID_RESOLUTION = 27
spatial_correlations = (
    atmos_flux_inversion.correlations.HomogeneousIsotropicCorrelation.
    from_function(
        atmos_flux_inversion.correlations.ExponentialCorrelation(
            CORRELATION_LENGTH / GRID_RESOLUTION),
        (len(TRUE_FLUXES_MATCHED.coords["dim_y"]),
         len(TRUE_FLUXES_MATCHED.coords["dim_x"])),
        is_cyclic=False))
# spatial_correlation_remapper = np.full(
#     # Grid points at full resolution, grid points at reduced resolution
#     (spatial_correlations.shape[0], 1),
#     # 1/Number of gridpoints combined in above mapping
#     1. / spatial_correlations.shape[0])
# # reduced_spatial_correlations = (
# #     spatial_correlation_remapper.dot(
# #         spatial_correlations.dot(spatial_correlation_remapper)))
import atmos_flux_inversion.remapper
spatial_obs_op_remapper, spatial_correlation_remapper = (
    atmos_flux_inversion.remapper.get_remappers(
        (len(TRUE_FLUXES_MATCHED.coords["dim_y"]),
         len(TRUE_FLUXES_MATCHED.coords["dim_x"])),
        UNCERTAINTY_RESOLUTION_REDUCTION_FACTOR))
REDUCED_NY, REDUCED_NX = spatial_correlation_remapper.shape[:2]
REDUCED_N_GRID_POINTS = REDUCED_NY * REDUCED_NX
write_progress_message("Have spatial correlations")

HOURLY_FLUX_TIMESCALE = 3
INTERVALS_PER_DAY = HOURS_PER_DAY // FLUX_INTERVAL
hour_correlations = (
    atmos_flux_inversion.correlations.HomogeneousIsotropicCorrelation.
    from_function(
        atmos_flux_inversion.correlations.ExponentialCorrelation(
            HOURLY_FLUX_TIMESCALE / FLUX_INTERVAL),
        (INTERVALS_PER_DAY,),
        is_cyclic=True))
hour_correlations_matrix = hour_correlations.dot(np.eye(
    hour_correlations.shape[0]))
write_progress_message("Have hourly correlations")
DAILY_FLUX_TIMESCALE = 7
DAILY_FLUX_FUN = "exp"
day_correlations = (
    atmos_flux_inversion.correlations.make_matrix(
        atmos_flux_inversion.correlations.ExponentialCorrelation(
            DAILY_FLUX_TIMESCALE
        ),
        (len(aligned_prior_fluxes.indexes["flux_time"]) *
         FLUX_INTERVAL // HOURS_PER_DAY,)))
write_progress_message("Have daily correlations")
temporal_correlations = kron(
    day_correlations,
    hour_correlations_matrix
).astype(np.float32)
print("Temporal:", type(temporal_correlations))
temporal_correlation_ds = xarray.DataArray(
    temporal_correlations,
    dict(flux_time=aligned_prior_fluxes.indexes["flux_time"].values,
         flux_time_adjoint=aligned_prior_fluxes.indexes["flux_time"].values),
    ("flux_time", "flux_time_adjoint"),
    "temporal_correlations",
    dict(long_name="temporal_correlations",
         units="dimensionless")
)
# Mean of fluxes, so sum of variances over square of number of
# elements in the group.  The number of rows and columns for each
# group will each be the number of elements in that group.
reduced_temporal_correlation_ds = (
    temporal_correlation_ds
    .resample(flux_time=UNCERTAINTY_TEMPORAL_RESOLUTION).mean("flux_time")
    .resample(flux_time_adjoint=UNCERTAINTY_TEMPORAL_RESOLUTION).mean(
        "flux_time_adjoint"
    )
)
write_progress_message("Have temporal correlations")
print(reduced_temporal_correlation_ds.values)
flush_output_streams()

full_correlations = kronecker_product(
    temporal_correlations,
    spatial_correlations)
print("Full:", type(full_correlations))
write_progress_message("Have combined correlations")

# I would like to add a fixed minimum at some point.
# full stds would then be sqrt(fixed^2 + varying^2)
# average seasonal variation (or some fraction thereof) might work.
# x2 since MsTMIP spread does not represent full uncertainty
# x5 since MsTMIP spread only represents monthly values and this uses sub-daily
# x10 matches model-model for Raczka for 200km/21d
# x4 matches model-model for 1000km/7d
FLUX_VARIANCE_VARYING_FRACTION = 4.
flux_std_pattern = xarray.open_dataset(
    "../data_files/2010_MsTMIP_flux_std.nc4",
    engine=NC_ENGINE
).get(
    ["E_TRA{:d}".format(i + 1) for i in range(1)]
).sel(
    Time=reduced_temporal_correlation_ds.indexes["flux_time"]
).mean(
    dim="Time",
    keep_attrs=True,
)

# Ensure units work out
for flux_part in flux_std_pattern.data_vars.values():
    unit = (cf_units.Unit(flux_part.attrs["units"]))
    if unit is not FLUX_UNITS:
        flux_part *= unit.convert(1, FLUX_UNITS)
        flux_part.attrs["units"] = str(FLUX_UNITS)

# flux_stds = (
#     FLUX_VARIANCE_VARYING_FRACTION *
#     flux_std_pattern["E_TRA1"].rolling(
#         center=True, min_periods=3, Time=21 * 8
#     ).mean(dim="Time").sel(
#         Time=temporal_correlation_ds.indexes["flux_time"]
#     ).data)
reduced_flux_stds = (
    FLUX_VARIANCE_VARYING_FRACTION *
    flux_std_pattern["E_TRA1"].data.astype(np.float32)
)
write_progress_message("Have standard deviations")

spatial_covariance = (
    atmos_flux_inversion.covariances.CorrelationStandardDeviation(
        spatial_correlations, reduced_flux_stds
    )
)
write_progress_message("Have full spatial covariance")

reduced_spatial_covariance = spatial_correlation_remapper.reshape(
    REDUCED_N_GRID_POINTS, N_GRID_POINTS
).dot(
    spatial_covariance.dot(
        spatial_correlation_remapper.reshape(
            REDUCED_N_GRID_POINTS, N_GRID_POINTS
        ).T
    )
)
write_progress_message("Have reduced spatial covariance")

write_progress_message("Have spatial covariances")
print(reduced_spatial_covariance)
flush_output_streams()

######################################################################
# Get an observation operator for the month
#
# To treat the flux as a mean, we need to sum the influence function
# spatial_obs_op_remapper_ds = xarray.DataArray(
#     spatial_obs_op_remapper,
#     dict(
#     ("reduced_dim_y", "reduced_dim_x", "dim_y", "dim_x"),
#     "spatial_obs_op_remapper_ds",
# )

reduced_influences = (
    aligned_influences
    .groupby_bins(
        "dim_x",
        pd.interval_range(
            0.,
            (aligned_influences.indexes["dim_x"][-1] +
             UNCERTAINTY_FLUX_RESOLUTION),
            freq=UNCERTAINTY_FLUX_RESOLUTION,
            closed="left")
        # np.arange(
        #     -1,
        #     (aligned_influences.indexes["dim_x"][-1] +
        #      UNCERTAINTY_FLUX_RESOLUTION),
        #     UNCERTAINTY_FLUX_RESOLUTION),
    ).sum("dim_x")
    .groupby_bins(
        "dim_y",
        pd.interval_range(
            0,
            (aligned_influences.indexes["dim_y"][-1] +
             UNCERTAINTY_FLUX_RESOLUTION),
            freq=UNCERTAINTY_FLUX_RESOLUTION, closed="left")
        # np.arange(
        #     -1,
        #      (aligned_influences.indexes["dim_y"][-1] +
        #       UNCERTAINTY_FLUX_RESOLUTION),
        #      UNCERTAINTY_FLUX_RESOLUTION),
    ).sum("dim_y")
    .resample(flux_time=UNCERTAINTY_TEMPORAL_RESOLUTION).sum("flux_time")
).rename(dim_x_bins="reduced_dim_x", dim_y_bins="reduced_dim_y",
         flux_time="reduced_flux_time")
sparse_reduced_influences_coords = sparse_influences.coords.copy()
# dim 0 is obs
# dim 1 is flux_time
assert UNCERTAINTY_TEMPORAL_RESOLUTION == "1W"
sparse_reduced_influences_coords[1, :] //= 7 * INTERVALS_PER_DAY
# dims 2 and 3 are y and x
sparse_reduced_influences_coords[2:, :] //= UNCERTAINTY_RESOLUTION_REDUCTION_FACTOR
reduced_influences_sparse = sparse.COO(
    sparse_reduced_influences_coords,
    sparse_influences.data.astype(np.float32)
)
del sparse_reduced_influences_coords

print(datetime.datetime.now(UTC).strftime("%c"),
      "Have reduced influence for covariance calculation.")
flush_output_streams()

prior_covariance = kronecker_product(
    temporal_correlations,
    spatial_covariance)
# Yay separability!
reduced_prior_covariance = kron(
    reduced_temporal_correlation_ds.data,
    reduced_spatial_covariance)
print("Covariance:", type(prior_covariance))
write_progress_message("Have covariances")
print(reduced_prior_covariance)
flush_output_streams()

# I realize this isn't quite the intended use for OBS_CHUNK
prior_fluxes = aligned_prior_fluxes.transpose(
    "flux_time", "dim_y", "dim_x", "realization")
write_progress_message("Have prior noise")

# TODO: use actual heights
here_obs = WRF_OBS_SITE[TRACER_NAME].sel(
    observation_time=xarray.DataArray(
        pd_obs_index,
        name="observation",
        dims=dict(observation=site_obs_pd_index)),
    site=xarray.DataArray(
        site_index,
        name="observation",
        dims=dict(observation=site_obs_pd_index)),
).rename(dict(projection_x_coordinate="tower_x",
              projection_y_coordinate="tower_y")).astype(np.float32)
print(here_obs)

OBSERVATION_STD = 2.
"""Standard deviation of observation transport error

This assumes similar deviations can be expected at each site.

Representativeness error from Gerbig et al 2003 for 27 km is .2 ppm
Ken says transport error is usually given as O(2-3ppm)
"""
OBS_CORR_FUN = atmos_flux_inversion.correlations.ExponentialCorrelation(3)
"""Temporal correlations in observation error.

Mostly reflects transport error.  Given in units of obs_time.
Since the input is hourly this is a three-hour decay time.
"""
OBS_INTERVAL = np.array(1, dtype='m8[h]')
observation_covariance = OBS_CORR_FUN(
    abs(pd_obs_index[:, np.newaxis] - pd_obs_index[np.newaxis, :]) /
    OBS_INTERVAL)
# Assumes no correlations between observations.
observation_covariance[
    site_index[:, np.newaxis] != site_index[np.newaxis, :]] = 0
observation_covariance *= OBSERVATION_STD ** 2
observation_covariance = asarray(observation_covariance).astype(np.float32)
obs_diag_index = np.arange(observation_covariance.shape[0])
# Add representativeness and instrument errors
observation_covariance[obs_diag_index, obs_diag_index] += 0.4 ** 2 + 0.1 ** 2

used_observation_vals = (
    here_obs.data[:, np.newaxis] +
    gaussian_noise(observation_covariance, N_REALIZATIONS).T.reshape(
        here_obs.shape + (N_REALIZATIONS,))).astype(np.float32).persist()
used_observations = xarray.DataArray(
    used_observation_vals,
    here_obs.coords,
    here_obs.dims + ("realization",),
    "pseudo_observations",
    dict(
        observation_standard_deviation=OBSERVATION_STD,
        observation_correlation_time=OBS_CORR_FUN._length)
).rename(dict(longitude_0="tower_lon", latitude_0="tower_lat"))
used_observations.coords["realization"] = range(N_REALIZATIONS)
used_observations.coords["realization"].attrs.update(dict(
    standard_name="realization"))
write_progress_message("Have observation noise")

print(datetime.datetime.now(UTC).strftime("%c"),
      "Got covariance parts, getting posterior")
flush_output_streams()
assert sparse_influences.shape == (
    aligned_influences.shape[0], N_FLUX_TIMES, NY, NX
)
assert reduced_influences_sparse.shape == reduced_influences.shape
posterior, reduced_posterior_covariances = (
    atmos_flux_inversion.optimal_interpolation.save_sum(
        aligned_prior_fluxes.values.reshape(
            N_GRID_POINTS * len(aligned_prior_fluxes.indexes["flux_time"]),
            N_REALIZATIONS).astype(np.float32),
        prior_covariance,
        used_observations.values.astype(np.float32),
        observation_covariance.astype(np.float32),
        sparse_influences.reshape((aligned_influences.shape[0],
                                   np.prod(aligned_influences.shape[-3:]))).astype(np.float32),
        reduced_prior_covariance,
        reduced_influences_sparse.reshape((reduced_influences.shape[0],
                                           np.prod(reduced_influences.shape[-3:]))).astype(np.float32)
    )
)
print(datetime.datetime.now(UTC).strftime("%c"),
      "Have posterior values, making dataset")
flush_output_streams()

posterior = posterior.reshape(aligned_prior_fluxes.shape)
posterior_ds = xarray.Dataset(
    dict(posterior=(aligned_prior_fluxes.dims, posterior,
                    posterior_var_atts),
         prior=(aligned_prior_fluxes.dims, aligned_prior_fluxes,
                aligned_prior_fluxes.attrs),
         truth=(aligned_true_fluxes.dims, aligned_true_fluxes,
                aligned_true_fluxes.attrs),
         ),
    aligned_prior_fluxes.coords,
    posterior_global_atts)
posterior_ds["pseudo_observations"] = used_observations
print(datetime.datetime.now(UTC).strftime("%c"),
      "Have posterior structure, evaluating and writing")
flush_output_streams()

encoding = {name: {"_FillValue": -99}
            for name in posterior_ds.data_vars}
encoding.update({name: {"_FillValue": None}
                 for name in posterior_ds.coords})

posterior_ds.to_netcdf(
    ("2010-07_monthly_inversion_{flux_interval:02d}h_027km_"
     "noise{ncorr_fun:s}{ncorr_len:d}km{ncorr_fun_time:s}{ncorr_len_time:d}d_"
     "icov{icorr_fun:s}{icorr_len:d}km{icorr_fun_time:s}{icorr_len_time:d}d_"
     "sparse_output_dense_spatial_corr.nc4")
    .format(flux_interval=FLUX_INTERVAL, ncorr_fun=CORR_FUN,
            ncorr_len=CORR_LEN, icorr_len=CORRELATION_LENGTH, icorr_fun="exp",
            icorr_len_time=DAILY_FLUX_TIMESCALE, icorr_fun_time=DAILY_FLUX_FUN,
            ncorr_fun_time=TIME_CORR_FUN, ncorr_len_time=TIME_CORR_LEN),
    encoding=encoding, engine=NC_ENGINE, mode="w")
write_progress_message("Wrote posterior")

write_progress_message(
    "Finding posterior covariance without corrections for "
    "aggregation error"
)

infl_fun_red = reduced_influences.stack(
    fluxes=("reduced_flux_time", "reduced_dim_y",
            "reduced_dim_x")
).values
B_HT_red = reduced_prior_covariance.dot(infl_fun_red.T)
red_post_cov_no_agg = reduced_prior_covariance - B_HT_red.dot(
    atmos_flux_inversion.linalg.solve(
        infl_fun_red.dot(B_HT_red) + observation_covariance,
        B_HT_red.T
    )
)
write_progress_message(
    "Found reduced posterior covariance without "
    "aggregation error approximations"
)

posterior_covariance_ds = xarray.Dataset(
    dict(
        reduced_posterior_covariance=(
            ("reduced_flux_time_adjoint",
             "reduced_dim_y_adjoint",
             "reduced_dim_x_adjoint",
             "reduced_flux_time",
             "reduced_dim_y",
             "reduced_dim_x",
             ),
            reduced_posterior_covariances.reshape(
                reduced_temporal_correlation_ds.shape[0],
                reduced_influences.shape[2],
                reduced_influences.shape[3],
                reduced_temporal_correlation_ds.shape[1],
                reduced_influences.shape[2],
                reduced_influences.shape[3]
            ),
            dict(
                standard_name=("surface_upward_mole_flux_of_carbon_dioxide "
                               "standard_error"),
                standard_error_multiplier=1.,
                long_name=(
                    "reduced_covariance_matrix_for_posterior_fluxes_full_HBHT"
                ),
                units=(FLUX_UNITS ** 2).format(),
                description=textwrap.dedent("""\
                Reduced-resolution approximation to the posterior
                covariance matrix, with no attempt to account for the
                increased aggregation error due to the increased
                resolution.
                """),
            ),
        ),
        reduced_posterior_covariance_no_aggregation=(
            ("reduced_flux_time_adjoint",
             "reduced_dim_y_adjoint",
             "reduced_dim_x_adjoint",
             "reduced_flux_time",
             "reduced_dim_y",
             "reduced_dim_x",
             ),
            red_post_cov_no_agg.reshape(
                reduced_temporal_correlation_ds.shape[0],
                reduced_influences.shape[2],
                reduced_influences.shape[3],
                reduced_temporal_correlation_ds.shape[1],
                reduced_influences.shape[2],
                reduced_influences.shape[3]
            ),
            dict(
                standard_name=("surface_upward_mole_flux_of_carbon_dioxide "
                               "standard_error"),
                standard_error_multiplier=1.,
                long_name=(
                    "reduced_covariance_matrix_for_posterior_fluxes_red_HBHT"
                ),
                units=(FLUX_UNITS ** 2).format(),
                description=textwrap.dedent("""\
                Reduced-resolution approximation to the posterior
                covariance matrix, with no attempt to account for the
                increased aggregation error due to the increased
                resolution.
                """),
            ),
        ),
        reduced_prior_covariance=(
            ("reduced_flux_time_adjoint",
             "reduced_dim_y_adjoint",
             "reduced_dim_x_adjoint",
             "reduced_flux_time",
             "reduced_dim_y",
             "reduced_dim_x",
             ),
            reduced_prior_covariance.reshape(
                reduced_temporal_correlation_ds.shape[0],
                reduced_influences.shape[2],
                reduced_influences.shape[3],
                reduced_temporal_correlation_ds.shape[1],
                reduced_influences.shape[2],
                reduced_influences.shape[3]
            ),
            dict(
                standard_name=("surface_upward_mole_flux_of_carbon_dioxide "
                               "standard_error"),
                standard_error_multiplier=1.,
                long_name="reduced_covariance_matrix_for_prior_fluxes",
                units=(FLUX_UNITS ** 2).format(),
            ),
        ),
    ),
    dict(
        reduced_flux_time=(
            reduced_temporal_correlation_ds.coords["flux_time"].values
        ),
        reduced_flux_time_adjoint=(
            reduced_temporal_correlation_ds.coords["flux_time_adjoint"].values
        ),
        wrf_proj=((), -1, WRF_PROJECTION.cf()),
    ),
    posterior_global_atts
)

RED_DIM_Y = reduced_influences.indexes["reduced_dim_y"]
posterior_covariance_ds.coords["reduced_dim_y"] = RED_DIM_Y.left
posterior_covariance_ds.coords["reduced_dim_y_bnds"] = (
    ("reduced_dim_y", "bnds2"),
    np.vstack([RED_DIM_Y.left, RED_DIM_Y.right]).T,
    dict(closed="left")
)
posterior_covariance_ds.coords["reduced_dim_y"].attrs.update(
    aligned_influences.coords["dim_y"].attrs)
posterior_covariance_ds.coords["reduced_dim_y"].attrs.update(
    dict(bounds="reduced_dim_y_bnds",
         adjoint="reduced_dim_y_adjoint"))

RED_DIM_X = reduced_influences.indexes["reduced_dim_x"]
posterior_covariance_ds.coords["reduced_dim_x"] = RED_DIM_X.left
posterior_covariance_ds.coords["reduced_dim_x_bnds"] = (
    ("reduced_dim_x", "bnds2"),
    np.vstack([RED_DIM_X.left, RED_DIM_X.right]).T,
    dict(closed="left"))
posterior_covariance_ds.coords["reduced_dim_x"].attrs.update(
    aligned_influences.coords["dim_x"].attrs)
posterior_covariance_ds.coords["reduced_dim_x"].attrs.update(
    dict(bounds="reduced_dim_x_bnds",
         adjoint="reduced_dim_x"))

posterior_covariance_ds.coords["reduced_dim_y_adjoint"] = (
    posterior_covariance_ds.coords["reduced_dim_y"].values)
posterior_covariance_ds.coords["reduced_dim_y_adjoint_bnds"] = (
    ("reduced_dim_y", "bnds2"),
    np.vstack([RED_DIM_Y.left, RED_DIM_Y.right]).T,
    dict(closed="left")
)
posterior_covariance_ds.coords["reduced_dim_y_adjoint"].attrs.update(
    aligned_influences.coords["dim_y"].attrs)
posterior_covariance_ds.coords["reduced_dim_y_adjoint"].attrs.update(
    dict(bounds="reduced_dim_y_adjoint_bnds",
         adjoint="reduced_dim_y"))

posterior_covariance_ds.coords["reduced_dim_x_adjoint"] = (
    posterior_covariance_ds.coords["reduced_dim_x"].values)
posterior_covariance_ds.coords["reduced_dim_x_adjoint_bnds"] = (
    ("reduced_dim_x", "bnds2"),
    np.vstack([RED_DIM_X.left, RED_DIM_X.right]).T,
    dict(closed="left"))
posterior_covariance_ds.attrs.update(
    dict(title="posterior_flux_uncertainty",
         summary="Covariance matrices for prior and posterior fluxes.",
         )
)

encoding = {name: {"_FillValue": -99}
            for name in posterior_covariance_ds.data_vars}
encoding.update({name: {"_FillValue": None}
                 for name in posterior_covariance_ds.coords})

posterior_covariance_ds.to_netcdf(
    ("2010-07_monthly_inversion_{flux_interval:02d}h_027km_"
     "noise{ncorr_fun:s}{ncorr_len:d}km{ncorr_fun_time:s}{ncorr_len_time:d}d_"
     "icov{icorr_fun:s}{icorr_len:d}km{icorr_fun_time:s}{icorr_len_time:d}d_"
     "{cov_res:d}km_{cov_tres:s}_covariance_output_sparse.nc4")
    .format(flux_interval=FLUX_INTERVAL, ncorr_fun=CORR_FUN,
            ncorr_len=CORR_LEN, icorr_len=CORRELATION_LENGTH, icorr_fun="exp",
            icorr_len_time=DAILY_FLUX_TIMESCALE, icorr_fun_time=DAILY_FLUX_FUN,
            ncorr_fun_time=TIME_CORR_FUN, ncorr_len_time=TIME_CORR_LEN,
            cov_res=int(UNCERTAINTY_FLUX_RESOLUTION // 1e3),
            cov_tres=UNCERTAINTY_TEMPORAL_RESOLUTION),
    encoding=encoding, engine=NC_ENGINE, mode="w")
write_progress_message("Wrote posterior covariance")
# UNCERTAINTY_FLUX_RESOLUTION UNCERTAINTY_TEMPORAL_RESOLUTION


text:
from .fcaide_ft_result import Save_result
import numpy as np
import math
import scipy.io as sio

class Fine_tuning:

    def __init__(self, save_file_name, weight_file_name, test_file_name, sigma_arr, model):

        self.save_file_name = save_file_name
        self.weight_file_name = weight_file_name
        self.test_file_name = test_file_name
        self.sigma_arr = sigma_arr

        self.init_lrate = 0.0001
        self.lrate_decay = 0.0
        self.mini_batch_size = 1
        self.max_ep = 30

        self.PSNR_arr_epoch = []
        self.SSIM_arr_epoch = []

        self.max_PSNR = 0
        self.max_SSIM = 0

        self.max_PSNR_denoised_image_arr = []
        self.max_SSIM_denoised_image_arr = []

        self.make_model = model

        return

    def get_images(self):

        f = sio.loadmat('./data/'+self.test_file_name)

        self.noisy_images = np.float32(np.array(f["noisy_images"]))
        self.clean_images = np.float32(np.array(f["clean_images"]))

        self.num_images = self.noisy_images.shape[0]
        self.img_x = self.noisy_images.shape[1]
        self.img_y = self.noisy_images.shape[2]

        self.noisy_images /= 255.
        self.clean_images /= 255.

        self.X = np.zeros((self.num_images,self.img_x,self.img_y,1))
        self.X[:,:,:,0] = self.noisy_images

        self.X -= 0.5
        self.X /= 0.2

        self.Y = np.zeros((self.num_images,self.img_x,self.img_y,3))
        self.Y[:,:,:,0] = self.clean_images
        self.Y[:,:,:,1] = self.noisy_images
        for idx in range(self.num_images):
            self.Y[idx,:,:,2] = self.sigma_arr[idx] / 255.

    def generate_flipped_image_set(self, X_data):

        if X_data.shape[2] == 1:

            flipped_image_set = []

            lr_flip = np.fliplr(X_data.reshape(self.img_x,self.img_y))
            ud_flip = np.flipud(X_data.reshape(self.img_x,self.img_y))
            lr_ud_flip = np.flipud(lr_flip)

            flipped_image_set = X_data.reshape(1,self.img_x,self.img_y,X_data.shape[2])
            flipped_image_set = np.vstack((flipped_image_set, lr_flip.reshape(1,self.img_x,self.img_y,X_data.shape[2])))
            flipped_image_set = np.vstack((flipped_image_set, ud_flip.reshape(1,self.img_x,self.img_y,X_data.shape[2])))
            flipped_image_set = np.vstack((flipped_image_set, lr_ud_flip.reshape(1,self.img_x,self.img_y,X_data.shape[2])))

        else:

            flipped_image_set = np.zeros((4,X_data.shape[0],X_data.shape[1],X_data.shape[2]))

            for i in range(3):

                origin = X_data[:,:,i]
                lr_flip = np.fliplr(X_data[:,:,i])
                ud_flip = np.flipud(X_data[:,:,i])
                lr_ud_flip = np.flipud(np.fliplr(X_data[:,:,i]))

                flipped_image_set[0,:,:,i] = origin
                flipped_image_set[1,:,:,i] = lr_flip
                flipped_image_set[2,:,:,i] = ud_flip
                flipped_image_set[3,:,:,i] = lr_ud_flip

        return flipped_image_set

    def test_images(self):

        self.get_images()
        save_results = Save_result(self.save_file_name, self.Y, self.img_x, self.img_y,)

        for idx in range(self.num_images):

            X_data = self.X[idx,:]
            X_data = self.generate_flipped_image_set(X_data)

            y_data = self.Y[idx,:]
            y_data = self.generate_flipped_image_set(y_data)

            X_data = X_data.reshape(4,self.img_x,self.img_y,1)
            y_data = y_data.reshape(4,self.img_x,self.img_y,3)

            self.model = self.make_model()
            self.model.load_weights('./weights/' + self.weight_file_name)

            self.model.fit(X_data, y_data, verbose=2, batch_size = self.mini_batch_size, epochs = self.max_ep, callbacks=[save_results])

            del self.model

reason: import numpy,import scipy


text:
import pandas as pd
import numpy as np

import sys
sys.path.append('.')

from spikelearn.data import SHORTCUTS

from numpy import dot
import matplotlib.pyplot as plt
import seaborn as sns

from itertools import product, count
from sklearn.preprocessing import StandardScaler
ss = StandardScaler()

# Load results data
data = pd.read_csv('data/results/across_trials/feature_importance.csv')
actv = pd.read_csv('data/results/across_trials/feature_importance_activity.csv')
score = pd.read_csv('data/results/across_trials/feature_importance_predictions.csv')

# WIll use a single train_size to compare scores
score = score.groupby('num_trials').get_group(50)

# Definitions used at generation
id_vars = ['rat','dataset','num_trials', 'unit', 'time','logC']

DRRD_RATS = SHORTCUTS['groups']['DRRD']
LOGCS = sorted(data.groupby('num_trials').get_group(50).logC.unique())[::-1]

# Preprocessing to get changes
each_mean = data.groupby(id_vars + ['when']).mean().reset_index(-1)
single_mean = data.groupby(id_vars).mean()
single_std = data.groupby(id_vars).std()

z_scored = (each_mean-single_mean)/single_std
z_scored['when'] = each_mean.when.values
data = z_scored.reset_index()

# Calculating score from predictions
score = score.groupby('num_trials').get_group(50)
score = pd.DataFrame(score.groupby(['rat', 'set', 'logC','cv','when']).apply(lambda x: x[['predictions','true']].corr().iloc[0,1]), columns=['score']).reset_index()

for rat, (fig_n, logC) in product(DRRD_RATS, enumerate(LOGCS)):
    one_rat = actv.groupby(['rat', 'logC', 'num_trials']).get_group((rat, logC, 50))
    fig = plt.figure(figsize=(12,24))
    for i, unit in enumerate(one_rat.unit.unique()):
        plt.subplot(one_rat.unit.nunique(), 3, 3*(i+1)-1)
        for when, c in [('init', 'b'),('end', 'r')]:
            local = one_rat.set_index(['unit','time']).groupby('when').get_group(when)
            local.loc[unit]['mean'].plot(color=c)
            plt.fill_between(local.loc[unit]['std'].index.values,
                             local.loc[unit]['mean'] + local.loc[0]['std'],
                             local.loc[unit]['mean'] - local.loc[0]['std'],alpha=.4,color=c)
            plt.axis('off');

    ax = plt.subplot(1,3,1)
    d = data.groupby(['rat', 'logC', 'num_trials','when']).get_group((rat, logC, 50, 'end'))
    d = d.pivot(index='unit', columns='time', values='value')
    sns.heatmap(d.values, ax=ax, vmin = -2, vmax = 2, cbar = False, cmap='RdBu_r')

    one_score = score.groupby(['rat','logC']).get_group((rat, logC)).groupby(['set','when']).mean()
    ax = plt.subplot(1,3,3)
    sns.barplot(x='set',y='score',hue='when',data=one_score.reset_index(), ax=ax,palette=['b','r'], hue_order=['init','end']); plt.ylim([0,1])
    plt.title('logC:%.3f'%logC, fontsize=30)

    plt.savefig('reports/figures/across_trials/feature_importance/feature_importance_{}_n{:3}_logC{:.3}.png'.format(rat, fig_n, logC), dpi=100)
    plt.close(fig)

max_stars_repo_name: EstevaoVieira/spikelearn

max_issues_repo_name: EstevaoVieira/spikelearn

max_forks_repo_name: EstevaoVieira/spikelearn

reason: import numpy,from numpy


text:
# Lint as: python3
#     Copyright 2020 Google LLC
#
#     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
#
#         https://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.
"""Wrapper for One-Class SVM Anomaly Detector based on sklearn."""
from madi.detectors.base_detector import BaseAnomalyDetectionAlgorithm
import madi.utils.sample_utils as sample_utils
import numpy as np
import pandas as pd
import sklearn.svm

class OneClassSVMAd(sklearn.svm.OneClassSVM, BaseAnomalyDetectionAlgorithm):
  """Wrapper class around the scikit-learn OC-SVM implementation."""

  def __init__(self, *args, **kwargs):
    """Constructs a OC-SVM Anomaly Detector.

    Args:
      *args: See the klearn.svm.OneClassSVM.
      **kwargs: See the klearn.svm.OneClassSVM.
    """
    super(OneClassSVMAd, self).__init__(*args, **kwargs)
    self._normalization_info = None

  def train_model(self, x_train: pd.DataFrame) -> None:
    """Trains a OC-SVM Anomaly detector using the positive sample.

    Args:
      x_train: training sample, with numeric feature columns.
    """
    self._normalization_info = sample_utils.get_normalization_info(x_train)
    column_order = sample_utils.get_column_order(self._normalization_info)
    normalized_x_train = sample_utils.normalize(x_train[column_order],
                                                self._normalization_info)
    super(OneClassSVMAd, self).fit(X=normalized_x_train)

  def predict(self, sample_df: pd.DataFrame) -> pd.DateOffset:
    """Performs anomaly detection on a new sample.

    Args:
      sample_df: dataframe with the new datapoints.

    Returns:
      original dataframe with a new column labled 'class_prob' as 1.0
      for normal to 0.0 for anomalous.
    """
    sample_df_normalized = sample_utils.normalize(sample_df,
                                                  self._normalization_info)
    column_order = sample_utils.get_column_order(self._normalization_info)
    x_test = np.float32(np.matrix(sample_df_normalized[column_order]))
    preds = super(OneClassSVMAd, self).predict(x_test)
    sample_df['class_prob'] = np.where(preds == -1, 0, preds)
    return sample_df

max_stars_repo_name: yllhaziri/madi

max_issues_repo_name: yllhaziri/madi

max_forks_repo_name: yllhaziri/madi


text:
# encoding:utf-8
# date:2020-10-26
# Author: X.L.Eric
# function: algo service

from flask import request, Flask, Response
import base64
import cv2
import numpy as np
import json
image_index = 0
app = Flask(__name__)

@app.route("/register", methods=['POST','GET'])
def get_frame():
    global image_index
    #解析图片数据
    img = base64.b64decode(str(request.form['image']))
    password = request.form["password"]
    user = request.form["user"]

    print("   user : {}, password : {}".format(user,password))

    image_data = np.fromstring(img, np.uint8)
    image_data = cv2.imdecode(image_data, cv2.IMREAD_COLOR)
    image_index += 1
    #------------------------ to do 算法解析

    cv2.imwrite('./service_get/example_{}.jpg'.format(image_index), image_data)
    print('{}) image_data shape : {} '.format(image_index,image_data.shape))

    dict_ = {
            "name":'test',
            }
    return Response(json.dumps(dict_),  mimetype='application/json')
if __name__ == "__main__":
    #多进程或多线程只能选择一个，不能同时开启
    # threaded=True
    # processes=True
    # 如果需要在通过flask 在不同计算机进行通信， host 需要设置为 0.0.0.0
    app.run(
        host = "localhost",
        port= 6666,
        debug = False,
        threaded = True,
        )


text:
import numpy as np
from numba import jit
import matplotlib.pyplot as plt
import seaborn as sns
import pandas as pd

def pairwise_distance(X):
    """
    Returns the pairwise distance matrix for a matrix X using negative squared Euclidean distance.
    """
    sum_X = (X ** 2).sum(axis=1)
    return -np.add(np.add(-2 * np.dot(X, X.T), sum_X).T, sum_X)

def prob_matrix(dist_X, sigmas):
    """
    Returns the matrix of conditional probabilities p_j|i.
    :param dist_X: the pairwise distance matrix
    :param sigmas: a vector of sigma values corresponding to each row of the distance matrix
    """
    two_sig_sq = np.asarray(2*np.square(sigmas)).reshape((-1, 1))
    x = dist_X / two_sig_sq
    expx = np.exp(x)

    # Since we do not sum over k = l, we set the diagonals to zero
    np.fill_diagonal(expx, 0.)
    # Avoid division by 0 errors
    expx = expx + 1e-8

    # Calculate Normalized Exponential
    rowsums = expx.sum(axis=1).reshape((-1, 1))
    normalized_exp = expx / rowsums
    return normalized_exp

def perplexity(prob_X):
    """
    Calculates the perplexity of each row of the probability matrix.
    :param prob_X: the conditional probability matrix
    :return: a vector of perplexity values
    """
    entropy = -np.sum(prob_X * np.log2(prob_X), axis=1)
    return 2**entropy

# add jit decorators
pairwise_distance_nb = jit(pairwise_distance, nopython=True, cache=True)
prob_matrix_nb = jit(prob_matrix)
perplexity_nb = jit(perplexity, nopython=True, cache=True)

def binary_search(f, target_perplexity, lower=1e-10, upper=1000, tol=1e-8, max_iter=10000):
    """
    Performs a binary search for the value of sigma that corresponds to the target perplexity.
    :param f: function to calculate perplexity
    :param target_perplexity: the specified perplexity value
    :param lower: initial lower bound
    :param upper: initial upper bound
    :param tol: tolerance to determine if the perplexity value is close enough to the target
    :param max_iter: maximum number of iterations of the loop
    :return: the optimal value of sigma
    """
    for i in range(max_iter):
        sigma = (lower + upper) / 2
        perp = f(sigma)
        if abs(perp - target_perplexity) < tol:
            return sigma
        if perp > target_perplexity:
            upper = sigma
        else:
            lower = sigma
    return sigma

def get_sigmas(dist_X, target_perplexity):
    """
    Finds the sigma for each row of the distance matrix based on the target perplexity.
    :param dist_X: the pairwise distance matrix
    :param target_perplexity: the specified perplexity value
    :return: a vector of sigma values corresponding to each row of the distance matrix
    """
    nrows = dist_X.shape[0]
    sigmas = np.zeros(nrows)

    for i in range(nrows):
        f = lambda sigma: perplexity_nb(prob_matrix_nb(dist_X[i:i + 1, :], np.asarray(sigma)))
        best_sigma = binary_search(f, target_perplexity)
        sigmas[i] = best_sigma
    return sigmas

def get_pmatrix(X, perplexity):
    """
    Calculates the final probability matrix using the pairwise affinities/conditional probabilities.
    :param M: the matrix of data to be converted
    :param perplexity: the specified perplexity
    :return: the joint probability matrix p_ij
    """
    # get the pairwise distances
    dist = pairwise_distance_nb(X)
    # get the sigmas
    sigmas = get_sigmas(dist, perplexity)
    # get the matrix of conditional probabilities
    prob = prob_matrix_nb(dist, sigmas)
    p = (prob + prob.T) / (2*prob.shape[0])
    return p

def get_qmatrix(Y):
    """
    Calculates the low dimensional affinities joint matrix q_ij.
    :param Y: low dimensional matrix representation of high dimensional matrix X
    :return: the joint probability matrix q_ij
    """
    q = 1 / (1 - pairwise_distance_nb(Y))
    np.fill_diagonal(q, 0)
    return q / q.sum()

def gradient(P, Q, Y):
    """
    Calculates the gradient of the Kullback-Leibler divergence between P and Q
    :param P: the joint probability matrix
    :param Q: the Student-t based joint probability distribution matrix
    :param Y: low dimensional matrix representation of high dimensional matrix X
    :return: a 2d array of the gradient values with the same dimensions as Y
    """
    pq_diff = np.expand_dims(P - Q, axis=2)
    y_diff = np.expand_dims(Y, axis=1) - np.expand_dims(Y, axis=0)
    dist = 1 - pairwise_distance_nb(Y)
    inv_distance = np.expand_dims(1 / dist, axis=2)

    # multiply and sum over each row
    grad = 4 * (pq_diff * y_diff * inv_distance).sum(axis=1)
    return grad

def TSNE(X, perplexity=40, num_iter=1000, learning_rate=100, momentum_initial=0.5, momentum_final=0.8):
    """
    Performs t-SNE on a given matrix X.
    :param X: matrix of high dimensional data
    :param perplexity: cost function parameter
    :param num_iter: number of iterations
    :param learning_rate: learning rate
    :param momentum_initial: initial momentum for first 250 iterations
    :param momentum_final: final momentum for remaining iterations
    :return: matrix of 2-dimensional data representing X
    """
    # calculate joint probability matrix
    joint_p = get_pmatrix(X, perplexity)
    # early exaggeration to improve optimization
    joint_p = 4 * joint_p

    # initialize Y by sampling from Gaussian
    Y_t = np.random.RandomState(1).normal(0, 10e-4, [X.shape[0], 2])
    # initialize past iteration Y_{t-1}
    Y_t1 = Y_t.copy()
    # initialize past iteration Y_{t-2}
    Y_t2 = Y_t.copy()

    for i in range(num_iter):
        # compute low dimensional affinities matrix
        joint_q = get_qmatrix(Y_t)
        # compute gradient
        grad = gradient(joint_p, joint_q, Y_t)

        # update momentum
        if i < 250:
            momentum = momentum_initial
        else:
            momentum = momentum_final

        # update current Y
        Y_t = Y_t1 - learning_rate * grad + momentum * (Y_t1 - Y_t2)
        # update past iterations
        Y_t1 = Y_t.copy()
        Y_t2 = Y_t1.copy()

        # conclude early exaggeration and revert joint probability matrix back
        if i == 50:
            joint_p = joint_p / 4

    return Y_t

def plot_tsne(Y, labels):
    """
    Plots a visualization of data processed through t-SNE
    :param Y: a two-dimensional matrix that has been reduced through t-SNE
    :param labels: an array of labels corresponding to each row of Y's classification
    :return: a plot of the data colored by class
    """
    df = pd.DataFrame(data=np.c_[Y, labels], columns=("Coord1", "Coord2", "label"))
    # split legend into two columns if there are lots of classification categories
    if len(set(labels)) > 10:
        sns.FacetGrid(df, hue="label", height=6).map(plt.scatter, "Coord1", "Coord2").add_legend(ncol=2)
    else:
        sns.FacetGrid(df, hue="label", height=6).map(plt.scatter, "Coord1", "Coord2").add_legend()
    plt.show()

max_stars_repo_name: sarahmansfield/t-SNE-py

max_issues_repo_name: sarahmansfield/t-SNE-py

max_forks_repo_name: sarahmansfield/t-SNE-py

reason: import numpy,from numba


text:
# -*- coding: utf-8 -*-
"""
pysteps.nowcasts.sseps
======================

Implementation of the Short-space ensemble prediction system (SSEPS) method.
Essentially, SSEPS is a localized version of STEPS.

For  localization  we  intend  the  use of  a  subset  of  the  observations
in  order  to  estimate  model parameters that are distributed in space.
The short-space approach used in :cite:`NBSG2017` is generalized to
the whole nowcasting system. This essenially boils down to a moving window
localization of the nowcasting procedure, whereby all parameters are estimated
over a subdomain of prescribed size.

.. autosummary::
    :toctree: ../generated/

    forecast
"""

import sys
import time

import numpy as np
import scipy.ndimage

from .. import cascade
from .. import extrapolation
from .. import noise
from ..nowcasts import utils as nowcast_utils
from ..postprocessing import probmatching
from ..timeseries import autoregression, correlation

try:
    import dask

    dask_imported = True
except ImportError:
    dask_imported = False

def forecast(
    R,
    metadata,
    V,
    n_timesteps,
    n_ens_members=24,
    n_cascade_levels=6,
    win_size=256,
    overlap=0.1,
    war_thr=0.1,
    extrap_method="semilagrangian",
    decomp_method="fft",
    bandpass_filter_method="gaussian",
    noise_method="ssft",
    ar_order=2,
    vel_pert_method=None,
    probmatching_method="cdf",
    mask_method="incremental",
    callback=None,
    fft_method="numpy",
    return_output=True,
    seed=None,
    num_workers=1,
    extrap_kwargs=None,
    filter_kwargs=None,
    noise_kwargs=None,
    vel_pert_kwargs=None,
    mask_kwargs=None,
    measure_time=False,
):
    """
    Generate a nowcast ensemble by using the Short-space ensemble prediction
    system (SSEPS) method.
    This is an experimental version of STEPS which allows for localization
    by means of a window function.

    Parameters
    ----------
    R : array-like
        Array of shape (ar_order+1,m,n) containing the input precipitation fields
        ordered by timestamp from oldest to newest. The time steps between the inputs
        are assumed to be regular, and the inputs are required to have finite values.
    metadata : dict
        Metadata dictionary containing the accutime, xpixelsize, threshold and
        zerovalue attributes as described in the documentation of
        :py:mod:`pysteps.io.importers`. xpixelsize is assumed to be in meters.
    V : array-like
        Array of shape (2,m,n) containing the x- and y-components of the advection
        field. The velocities are assumed to represent one time step between the
        inputs. All values are required to be finite.
    win_size : int or two-element sequence of ints
        Size-length of the localization window.
    overlap : float [0,1[
        A float between 0 and 1 prescribing the level of overlap between
        successive windows. If set to 0, no overlap is used.
    war_thr : float
        Threshold for the minimum fraction of rain in a given window.
    n_timesteps : int
        Number of time steps to forecast.
    n_ens_members : int
        The number of ensemble members to generate.
    n_cascade_levels : int
        The number of cascade levels to use.

    extrap_method : {'semilagrangian'}
        Name of the extrapolation method to use. See the documentation of
        pysteps.extrapolation.interface.
    decomp_method : {'fft'}
        Name of the cascade decomposition method to use. See the documentation
        of pysteps.cascade.interface.
    bandpass_filter_method : {'gaussian', 'uniform'}
        Name of the bandpass filter method to use with the cascade
        decomposition.
    noise_method : {'parametric','nonparametric','ssft','nested',None}
        Name of the noise generator to use for perturbating the precipitation
        field. See the documentation of pysteps.noise.interface. If set to None,
        no noise is generated.
    ar_order: int
        The order of the autoregressive model to use. Must be >= 1.
    vel_pert_method: {'bps',None}
        Name of the noise generator to use for perturbing the advection field.
        See the documentation of pysteps.noise.interface. If set to None,
        the advection field is not perturbed.
    mask_method : {'incremental', None}
        The method to use for masking no precipitation areas in the forecast
        field. The masked pixels are set to the minimum value of the
        observations. 'incremental' = iteratively buffer the mask with a
        certain rate (currently it is 1 km/min), None=no masking.
    probmatching_method : {'cdf', None}
        Method for matching the statistics of the forecast field with those of
        the most recently observed one. 'cdf'=map the forecast CDF to the
        observed one, None=no matching applied. Using 'mean' requires
        that mask_method is not None.
    callback : function
        Optional function that is called after computation of each time step of
        the nowcast. The function takes one argument: a three-dimensional array
        of shape (n_ens_members,h,w), where h and w are the height and width
        of the input field R, respectively. This can be used, for instance,
        writing the outputs into files.
    return_output : bool
        Set to False to disable returning the outputs as numpy arrays. This can
        save memory if the intermediate results are written to output files
        using the callback function.
    seed : int
        Optional seed number for the random generators.
    num_workers : int
        The number of workers to use for parallel computation. Applicable if
        dask is enabled or pyFFTW is used for computing the FFT.
        When num_workers>1, it is advisable to disable OpenMP by setting the
        environment variable OMP_NUM_THREADS to 1.
        This avoids slowdown caused by too many simultaneous threads.
    fft_method : str
        A string defining the FFT method to use (see utils.fft.get_method).
        Defaults to 'numpy' for compatibility reasons. If pyFFTW is installed,
        the recommended method is 'pyfftw'.
    extrap_kwargs : dict
        Optional dictionary containing keyword arguments for the extrapolation
        method. See the documentation of pysteps.extrapolation.
    filter_kwargs : dict
        Optional dictionary containing keyword arguments for the filter method.
        See the documentation of pysteps.cascade.bandpass_filters.py.
    noise_kwargs : dict
        Optional dictionary containing keyword arguments for the initializer of
        the noise generator. See the documentation of
        pysteps.noise.fftgenerators.
    vel_pert_kwargs : dict
        Optional dictionary containing keyword arguments "p_pert_par" and
        "p_pert_perp" for the initializer of the velocity perturbator.
        See the documentation of pysteps.noise.motion.
    mask_kwargs : dict
        Optional dictionary containing mask keyword arguments 'mask_f' and
        'mask_rim', the factor defining the the mask increment and the rim size,
        respectively.
        The mask increment is defined as mask_f*timestep/kmperpixel.
    measure_time : bool
        If set to True, measure, print and return the computation time.

    Returns
    -------
    out : ndarray
        If return_output is True, a four-dimensional array of shape
        (n_ens_members,n_timesteps,m,n) containing a time series of forecast
        precipitation fields for each ensemble member. Otherwise, a None value
        is returned. The time series starts from t0+timestep, where timestep is
        taken from the input precipitation fields R.

    See also
    --------
    pysteps.extrapolation.interface, pysteps.cascade.interface,
    pysteps.noise.interface, pysteps.noise.utils.compute_noise_stddev_adjs

    Notes
    -----
    Please be aware that this represents a (very) experimental implementation.

    References
    ----------
    :cite:`Seed2003`, :cite:`BPS2006`, :cite:`SPN2013`, :cite:`NBSG2017`

    """
    _check_inputs(R, V, ar_order)

    if extrap_kwargs is None:
        extrap_kwargs = dict()

    if filter_kwargs is None:
        filter_kwargs = dict()

    if noise_kwargs is None:
        noise_kwargs = dict()

    if vel_pert_kwargs is None:
        vel_pert_kwargs = dict()

    if mask_kwargs is None:
        mask_kwargs = dict()

    if np.any(~np.isfinite(R)):
        raise ValueError("R contains non-finite values")

    if np.any(~np.isfinite(V)):
        raise ValueError("V contains non-finite values")

    if mask_method not in ["incremental", None]:
        raise ValueError(
            "unknown mask method %s: must be 'incremental' or None" % mask_method
        )

    if np.isscalar(win_size):
        win_size = (np.int(win_size), np.int(win_size))
    else:
        win_size = tuple([np.int(win_size[i]) for i in range(2)])

    timestep = metadata["accutime"]
    kmperpixel = metadata["xpixelsize"] / 1000

    print("Computing SSEPS nowcast:")
    print("------------------------")
    print("")

    print("Inputs:")
    print("-------")
    print("input dimensions: %dx%d" % (R.shape[1], R.shape[2]))
    print("km/pixel:         %g" % kmperpixel)
    print("time step:        %d minutes" % timestep)
    print("")

    print("Methods:")
    print("--------")
    print("extrapolation:          %s" % extrap_method)
    print("bandpass filter:        %s" % bandpass_filter_method)
    print("decomposition:          %s" % decomp_method)
    print("noise generator:        %s" % noise_method)
    print("velocity perturbator:   %s" % vel_pert_method)
    print("precip. mask method:    %s" % mask_method)
    print("probability matching:   %s" % probmatching_method)
    print("FFT method:             %s" % fft_method)
    print("")

    print("Parameters:")
    print("-----------")
    print("localization window:      %dx%d" % (win_size[0], win_size[1]))
    print("overlap:                  %.1f" % overlap)
    print("war thr:                  %.2f" % war_thr)
    print("number of time steps:     %d" % n_timesteps)
    print("ensemble size:            %d" % n_ens_members)
    print("number of cascade levels: %d" % n_cascade_levels)
    print("order of the AR(p) model: %d" % ar_order)
    print("dask imported:            %s" % ("yes" if dask_imported else "no"))
    print("num workers:              %d" % num_workers)

    if vel_pert_method is "bps":
        vp_par = vel_pert_kwargs.get(
            "p_pert_par", noise.motion.get_default_params_bps_par()
        )
        vp_perp = vel_pert_kwargs.get(
            "p_pert_perp", noise.motion.get_default_params_bps_perp()
        )
        print(
            "velocity perturbations, parallel:      %g,%g,%g"
            % (vp_par[0], vp_par[1], vp_par[2])
        )
        print(
            "velocity perturbations, perpendicular: %g,%g,%g"
            % (vp_perp[0], vp_perp[1], vp_perp[2])
        )

    R_thr = metadata["threshold"]
    R_min = metadata["zerovalue"]

    num_ensemble_workers = n_ens_members if num_workers > n_ens_members else num_workers

    if measure_time:
        starttime_init = time.time()

    # get methods
    extrapolator_method = extrapolation.get_method(extrap_method)

    x_values, y_values = np.meshgrid(np.arange(R.shape[2]), np.arange(R.shape[1]))

    xy_coords = np.stack([x_values, y_values])

    decomp_method, __ = cascade.get_method(decomp_method)
    filter_method = cascade.get_method(bandpass_filter_method)
    if noise_method is not None:
        init_noise, generate_noise = noise.get_method(noise_method)

    # advect the previous precipitation fields to the same position with the
    # most recent one (i.e. transform them into the Lagrangian coordinates)
    R = R[-(ar_order + 1) :, :, :].copy()
    extrap_kwargs = extrap_kwargs.copy()
    extrap_kwargs["xy_coords"] = xy_coords
    res = []
    f = lambda R, i: extrapolator_method(
        R[i, :, :], V, ar_order - i, "min", **extrap_kwargs
    )[-1]
    for i in range(ar_order):
        if not dask_imported:
            R[i, :, :] = f(R, i)
        else:
            res.append(dask.delayed(f)(R, i))

    if dask_imported:
        num_workers_ = len(res) if num_workers > len(res) else num_workers
        R = np.stack(list(dask.compute(*res, num_workers=num_workers_)) + [R[-1, :, :]])

    if mask_method == "incremental":
        # get mask parameters
        mask_rim = mask_kwargs.get("mask_rim", 10)
        mask_f = mask_kwargs.get("mask_f", 1.0)
        # initialize the structuring element
        struct = scipy.ndimage.generate_binary_structure(2, 1)
        # iterate it to expand it nxn
        n = mask_f * timestep / kmperpixel
        struct = scipy.ndimage.iterate_structure(struct, int((n - 1) / 2.0))

    noise_kwargs.update(
        {
            "win_size": win_size,
            "overlap": overlap,
            "war_thr": war_thr,
            "rm_rdisc": True,
            "donorm": True,
        }
    )

    print("Estimating nowcast parameters...", end="")

    def estimator(R, parsglob=None, idxm=None, idxn=None):

        pars = {}

        # initialize the perturbation generator for the precipitation field
        if noise_method is not None and parsglob is None:
            P = init_noise(R, fft_method=fft_method, **noise_kwargs)
        else:
            P = None
        pars["P"] = P

        # initialize the band-pass filter
        if parsglob is None:
            filter = filter_method(R.shape[1:], n_cascade_levels, **filter_kwargs)
            pars["filter"] = filter
        else:
            pars["filter"] = None

        # compute the cascade decompositions of the input precipitation fields
        if parsglob is None:
            R_d = []
            for i in range(ar_order + 1):
                R_d_ = decomp_method(R[i, :, :], filter, fft_method=fft_method,
                                     normalize=True, compute_stats=True)
                R_d.append(R_d_)
            R_d_ = None

        # normalize the cascades and rearrange them into a four-dimensional array
        # of shape (n_cascade_levels,ar_order+1,m,n) for the autoregressive model
        if parsglob is None:
            R_c = nowcast_utils.stack_cascades(R_d, n_cascade_levels)
            mu = R_d[-1]["means"]
            sigma = R_d[-1]["stds"]
            R_d = None

        else:
            R_c = parsglob["R_c"][0][
                :, :, idxm.item(0) : idxm.item(1), idxn.item(0) : idxn.item(1)
            ].copy()
            mu = np.mean(R_c, axis=(2, 3))
            sigma = np.std(R_c, axis=(2, 3))

            R_c = (R_c - mu[:, :, None, None]) / sigma[:, :, None, None]

            mu = mu[:, -1]
            sigma = sigma[:, -1]

        pars["mu"] = mu
        pars["sigma"] = sigma

        # compute lag-l temporal autocorrelation coefficients for each cascade level
        GAMMA = np.empty((n_cascade_levels, ar_order))
        for i in range(n_cascade_levels):
            R_c_ = np.stack([R_c[i, j, :, :] for j in range(ar_order + 1)])
            GAMMA[i, :] = correlation.temporal_autocorrelation(R_c_)
        R_c_ = None

        if ar_order == 2:
            # adjust the local lag-2 correlation coefficient to ensure that the AR(p)
            # process is stationary
            for i in range(n_cascade_levels):
                GAMMA[i, 1] = autoregression.adjust_lag2_corrcoef2(
                    GAMMA[i, 0], GAMMA[i, 1]
                )

        # estimate the parameters of the AR(p) model from the autocorrelation
        # coefficients
        PHI = np.empty((n_cascade_levels, ar_order + 1))
        for i in range(n_cascade_levels):
            PHI[i, :] = autoregression.estimate_ar_params_yw(GAMMA[i, :])
        pars["PHI"] = PHI

        # stack the cascades into a five-dimensional array containing all ensemble
        # members
        R_c = [R_c.copy() for i in range(n_ens_members)]
        pars["R_c"] = R_c

        if mask_method is not None and parsglob is None:
            MASK_prec = R[-1, :, :] >= R_thr
            if mask_method == "incremental":
                # initialize precip mask for each member
                MASK_prec = _compute_incremental_mask(MASK_prec, struct, mask_rim)
                MASK_prec = [MASK_prec.copy() for j in range(n_ens_members)]
        else:
            MASK_prec = None
        pars["MASK_prec"] = MASK_prec

        return pars

    # prepare windows
    M, N = R.shape[1:]
    n_windows_M = np.ceil(1.0 * M / win_size[0]).astype(int)
    n_windows_N = np.ceil(1.0 * N / win_size[1]).astype(int)
    idxm = np.zeros((2, 1), dtype=int)
    idxn = np.zeros((2, 1), dtype=int)

    sys.stdout.flush()
    if measure_time:
        starttime = time.time()

    # compute global parameters to be used as defaults
    parsglob = estimator(R)

    # loop windows
    if n_windows_M > 1 or n_windows_N > 1:
        war = np.empty((n_windows_M, n_windows_N))
        PHI = np.empty((n_windows_M, n_windows_N, n_cascade_levels, ar_order + 1))
        mu = np.empty((n_windows_M, n_windows_N, n_cascade_levels))
        sigma = np.empty((n_windows_M, n_windows_N, n_cascade_levels))
        ff = []
        rc = []
        pp = []
        mm = []
        for m in range(n_windows_M):
            ff_ = []
            pp_ = []
            rc_ = []
            mm_ = []
            for n in range(n_windows_N):

                # compute indices of local window
                idxm[0] = int(np.max((m * win_size[0] - overlap * win_size[0], 0)))
                idxm[1] = int(
                    np.min((idxm[0] + win_size[0] + overlap * win_size[0], M))
                )
                idxn[0] = int(np.max((n * win_size[1] - overlap * win_size[1], 0)))
                idxn[1] = int(
                    np.min((idxn[0] + win_size[1] + overlap * win_size[1], N))
                )

                mask = np.zeros((M, N), dtype=bool)
                mask[idxm.item(0) : idxm.item(1), idxn.item(0) : idxn.item(1)] = True

                R_ = R[:, idxm.item(0) : idxm.item(1), idxn.item(0) : idxn.item(1)]

                war[m, n] = np.sum(R_[-1, :, :] >= R_thr) / R_[-1, :, :].size
                if war[m, n] > war_thr:

                    # estimate local parameters
                    pars = estimator(R, parsglob, idxm, idxn)
                    ff_.append(pars["filter"])
                    pp_.append(pars["P"])
                    rc_.append(pars["R_c"])
                    mm_.append(pars["MASK_prec"])
                    mu[m, n, :] = pars["mu"]
                    sigma[m, n, :] = pars["sigma"]
                    PHI[m, n, :, :] = pars["PHI"]

                else:
                    # dry window
                    ff_.append(None)
                    pp_.append(None)
                    rc_.append(None)
                    mm_.append(None)

            ff.append(ff_)
            pp.append(pp_)
            rc.append(rc_)
            mm.append(mm_)

        # remove unnecessary variables
        ff_ = None
        pp_ = None
        rc_ = None
        mm_ = None
        pars = None

    if measure_time:
        print("%.2f seconds." % (time.time() - starttime))
    else:
        print(" done.")

    # initialize the random generators
    if noise_method is not None:
        randgen_prec = []
        randgen_motion = []
        np.random.seed(seed)
        for j in range(n_ens_members):
            rs = np.random.RandomState(seed)
            randgen_prec.append(rs)
            seed = rs.randint(0, high=1e9)
            rs = np.random.RandomState(seed)
            randgen_motion.append(rs)
            seed = rs.randint(0, high=1e9)

    if vel_pert_method is not None:
        init_vel_noise, generate_vel_noise = noise.get_method(vel_pert_method)

        # initialize the perturbation generators for the motion field
        vps = []
        for j in range(n_ens_members):
            kwargs = {
                "randstate": randgen_motion[j],
                "p_par": vp_par,
                "p_perp": vp_perp,
            }
            vp_ = init_vel_noise(V, 1.0 / kmperpixel, timestep, **kwargs)
            vps.append(vp_)

    D = [None for j in range(n_ens_members)]
    R_f = [[] for j in range(n_ens_members)]

    if measure_time:
        init_time = time.time() - starttime_init

    R = R[-1, :, :]

    print("Starting nowcast computation.")

    if measure_time:
        starttime_mainloop = time.time()

    # iterate each time step
    for t in range(n_timesteps):
        print("Computing nowcast for time step %d... " % (t + 1), end="")
        sys.stdout.flush()
        if measure_time:
            starttime = time.time()

        # iterate each ensemble member
        def worker(j):

            # first the global step

            if noise_method is not None:
                # generate noise field
                EPS = generate_noise(
                    parsglob["P"], randstate=randgen_prec[j], fft_method=fft_method
                )
                # decompose the noise field into a cascade
                EPS_d = decomp_method(EPS, parsglob["filter"], fft_method=fft_method,
                                      normalize=True, compute_stats=True)
            else:
                EPS_d = None

            # iterate the AR(p) model for each cascade level
            R_c = parsglob["R_c"][j].copy()
            if R_c.shape[1] >= ar_order:
                R_c = R_c[:, -ar_order:, :, :].copy()
            for i in range(n_cascade_levels):
                # normalize the noise cascade
                if EPS_d is not None:
                    EPS_ = (
                        EPS_d["cascade_levels"][i, :, :] - EPS_d["means"][i]
                    ) / EPS_d["stds"][i]
                else:
                    EPS_ = None
                # apply AR(p) process to cascade level
                R_c[i, :, :, :] = autoregression.iterate_ar_model(
                    R_c[i, :, :, :], parsglob["PHI"][i, :], eps=EPS_
                )
                EPS_ = None
            parsglob["R_c"][j] = R_c.copy()
            EPS = None

            # compute the recomposed precipitation field(s) from the cascades
            # obtained from the AR(p) model(s)
            R_c_ = _recompose_cascade(R_c, parsglob["mu"], parsglob["sigma"])
            R_c = None

            # then the local steps
            if n_windows_M > 1 or n_windows_N > 1:
                idxm = np.zeros((2, 1), dtype=int)
                idxn = np.zeros((2, 1), dtype=int)
                R_l = np.zeros((M, N), dtype=float)
                M_s = np.zeros((M, N), dtype=float)
                for m in range(n_windows_M):
                    for n in range(n_windows_N):

                        # compute indices of local window
                        idxm[0] = int(
                            np.max((m * win_size[0] - overlap * win_size[0], 0))
                        )
                        idxm[1] = int(
                            np.min((idxm[0] + win_size[0] + overlap * win_size[0], M))
                        )
                        idxn[0] = int(
                            np.max((n * win_size[1] - overlap * win_size[1], 0))
                        )
                        idxn[1] = int(
                            np.min((idxn[0] + win_size[1] + overlap * win_size[1], N))
                        )

                        # build localization mask
                        mask = _get_mask((M, N), idxm, idxn)
                        mask_l = mask[
                            idxm.item(0) : idxm.item(1), idxn.item(0) : idxn.item(1)
                        ]
                        M_s += mask

                        # skip if dry
                        if war[m, n] > war_thr:

                            R_c = rc[m][n][j].copy()
                            if R_c.shape[1] >= ar_order:
                                R_c = R_c[:, -ar_order:, :, :]
                            if noise_method is not None:
                                # extract noise field
                                EPS_d_l = EPS_d["cascade_levels"][
                                    :,
                                    idxm.item(0) : idxm.item(1),
                                    idxn.item(0) : idxn.item(1),
                                ].copy()
                                mu_ = np.mean(EPS_d_l, axis=(1, 2))
                                sigma_ = np.std(EPS_d_l, axis=(1, 2))
                            else:
                                EPS_d_l = None

                            # iterate the AR(p) model for each cascade level
                            for i in range(n_cascade_levels):
                                # normalize the noise cascade
                                if EPS_d_l is not None:
                                    EPS_ = (
                                        EPS_d_l[i, :, :] - mu_[i, None, None]
                                    ) / sigma_[i, None, None]
                                else:
                                    EPS_ = None
                                # apply AR(p) process to cascade level
                                R_c[i, :, :, :] = autoregression.iterate_ar_model(
                                    R_c[i, :, :, :], PHI[m, n, i, :], eps=EPS_
                                )
                                EPS_ = None
                            rc[m][n][j] = R_c.copy()
                            EPS_d_l = mu_ = sigma_ = None

                            # compute the recomposed precipitation field(s) from the cascades
                            # obtained from the AR(p) model(s)
                            mu_ = mu[m, n, :]
                            sigma_ = sigma[m, n, :]
                            R_c = [
                                ((R_c[i, -1, :, :] * sigma_[i]) + mu_[i])
                                * parsglob["sigma"][i]
                                + parsglob["mu"][i]
                                for i in range(len(mu_))
                            ]
                            R_l_ = np.sum(np.stack(R_c), axis=0)
                            R_c = mu_ = sigma_ = None
                            # R_l_ = _recompose_cascade(R_c[:, :, :], mu[m, n, :], sigma[m, n, :])
                        else:
                            R_l_ = R_c_[
                                idxm.item(0) : idxm.item(1), idxn.item(0) : idxn.item(1)
                            ].copy()

                        if probmatching_method == "cdf":
                            # adjust the CDF of the forecast to match the most recently
                            # observed precipitation field
                            R_ = R[
                                idxm.item(0) : idxm.item(1), idxn.item(0) : idxn.item(1)
                            ].copy()
                            R_l_ = probmatching.nonparam_match_empirical_cdf(R_l_, R_)
                            R_ = None

                        R_l[
                            idxm.item(0) : idxm.item(1), idxn.item(0) : idxn.item(1)
                        ] += (R_l_ * mask_l)
                        R_l_ = None

                ind = M_s > 0
                R_l[ind] *= 1 / M_s[ind]
                R_l[~ind] = R_min

                R_c_ = R_l.copy()
                R_l = None

            if probmatching_method == "cdf":
                # adjust the CDF of the forecast to match the most recently
                # observed precipitation field
                R_c_[R_c_ < R_thr] = R_min
                R_c_ = probmatching.nonparam_match_empirical_cdf(R_c_, R)

            if mask_method is not None:
                # apply the precipitation mask to prevent generation of new
                # precipitation into areas where it was not originally
                # observed
                if mask_method == "incremental":
                    MASK_prec = parsglob["MASK_prec"][j].copy()
                    R_c_ = R_c_.min() + (R_c_ - R_c_.min()) * MASK_prec
                    MASK_prec = None

            if mask_method == "incremental":
                parsglob["MASK_prec"][j] = _compute_incremental_mask(
                    R_c_ >= R_thr, struct, mask_rim
                )

            # compute the perturbed motion field
            if vel_pert_method is not None:
                V_ = V + generate_vel_noise(vps[j], (t + 1) * timestep)
            else:
                V_ = V

            # advect the recomposed precipitation field to obtain the forecast
            # for time step t
            extrap_kwargs.update({"D_prev": D[j], "return_displacement": True})
            R_f_, D_ = extrapolator_method(R_c_, V_, 1, **extrap_kwargs)
            D[j] = D_
            R_f_ = R_f_[0]

            R_f_[R_f_ < R_thr] = R_min

            return R_f_

        res = []
        for j in range(n_ens_members):
            if not dask_imported or n_ens_members == 1:
                res.append(worker(j))
            else:
                res.append(dask.delayed(worker)(j))

        R_f_ = (
            dask.compute(*res, num_workers=num_ensemble_workers)
            if dask_imported and n_ens_members > 1
            else res
        )
        res = None

        if measure_time:
            print("%.2f seconds." % (time.time() - starttime))
        else:
            print("done.")

        if callback is not None:
            callback(np.stack(R_f_))
            R_f_ = None

        if return_output:
            for j in range(n_ens_members):
                R_f[j].append(R_f_[j])

    if measure_time:
        mainloop_time = time.time() - starttime_mainloop

    if return_output:
        outarr = np.stack([np.stack(R_f[j]) for j in range(n_ens_members)])
        if measure_time:
            return outarr, init_time, mainloop_time
        else:
            return outarr
    else:
        return None

def _check_inputs(R, V, ar_order):
    if len(R.shape) != 3:
        raise ValueError("R must be a three-dimensional array")
    if R.shape[0] < ar_order + 1:
        raise ValueError("R.shape[0] < ar_order+1")
    if len(V.shape) != 3:
        raise ValueError("V must be a three-dimensional array")
    if R.shape[1:3] != V.shape[1:3]:
        raise ValueError(
            "dimension mismatch between R and V: shape(R)=%s, shape(V)=%s"
            % (str(R.shape), str(V.shape))
        )

def _compute_incremental_mask(Rbin, kr, r):
    # buffer the observation mask Rbin using the kernel kr
    # add a grayscale rim r (for smooth rain/no-rain transition)

    # buffer observation mask
    Rbin = np.ndarray.astype(Rbin.copy(), "uint8")
    Rd = scipy.ndimage.morphology.binary_dilation(Rbin, kr)

    # add grayscale rim
    kr1 = scipy.ndimage.generate_binary_structure(2, 1)
    mask = Rd.astype(int)
    for n in range(r):
        Rd = scipy.ndimage.morphology.binary_dilation(Rd, kr1)
        mask += Rd
    # normalize between 0 and 1
    return mask / mask.max()

# TODO: Use the recomponse_cascade method in the cascade.decomposition module
def _recompose_cascade(R, mu, sigma):
    R_rc = [(R[i, -1, :, :] * sigma[i]) + mu[i] for i in range(len(mu))]
    R_rc = np.sum(np.stack(R_rc), axis=0)

    return R_rc

def _build_2D_tapering_function(win_size, win_type="flat-hanning"):
    """Produces two-dimensional tapering function for rectangular fields.

    Parameters
    ----------
    win_size : tuple of int
        Size of the tapering window as two-element tuple of integers.
    win_type : str
        Name of the tapering window type (hanning, flat-hanning)

    Returns
    -------
    w2d : array-like
        A two-dimensional numpy array containing the 2D tapering function.
    """

    if len(win_size) != 2:
        raise ValueError("win_size is not a two-element tuple")

    if win_type == "hanning":
        w1dr = np.hanning(win_size[0])
        w1dc = np.hanning(win_size[1])

    elif win_type == "flat-hanning":

        T = win_size[0] / 4.0
        W = win_size[0] / 2.0
        B = np.linspace(-W, W, int(2 * W))
        R = np.abs(B) - T
        R[R < 0] = 0.0
        A = 0.5 * (1.0 + np.cos(np.pi * R / T))
        A[np.abs(B) > (2 * T)] = 0.0
        w1dr = A

        T = win_size[1] / 4.0
        W = win_size[1] / 2.0
        B = np.linspace(-W, W, int(2 * W))
        R = np.abs(B) - T
        R[R < 0] = 0.0
        A = 0.5 * (1.0 + np.cos(np.pi * R / T))
        A[np.abs(B) > (2 * T)] = 0.0
        w1dc = A

    elif win_type == "rectangular":

        w1dr = np.ones(win_size[0])
        w1dc = np.ones(win_size[1])

    else:
        raise ValueError("unknown win_type %s" % win_type)

    # Expand to 2-D
    # w2d = np.sqrt(np.outer(w1dr,w1dc))
    w2d = np.outer(w1dr, w1dc)

    # Set nans to zero
    if np.any(np.isnan(w2d)):
        w2d[np.isnan(w2d)] = np.min(w2d[w2d > 0])

    w2d[w2d < 1e-3] = 1e-3

    return w2d

def _get_mask(Size, idxi, idxj, win_type="flat-hanning"):
    """Compute a mask of zeros with a window at a given position.
    """

    idxi = np.array(idxi).astype(int)
    idxj = np.array(idxj).astype(int)

    win_size = (idxi[1] - idxi[0], idxj[1] - idxj[0])
    wind = _build_2D_tapering_function(win_size, win_type)

    mask = np.zeros(Size)
    mask[idxi.item(0) : idxi.item(1), idxj.item(0) : idxj.item(1)] = wind

    return mask

max_stars_repo_name: alexanderhucheerful/pysteps

max_issues_repo_name: alexanderhucheerful/pysteps

max_forks_repo_name: alexanderhucheerful/pysteps


text:
# Copyright 2019 Google LLC
#
# 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
#
#     https://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.

import os
import numpy as np
from PIL import Image
from PIL import ImageDraw
from pose_engine import PoseEngine
import svgwrite

EDGES = (
    ('nose', 'left eye'),
    ('nose', 'right eye'),
    ('nose', 'left ear'),
    ('nose', 'right ear'),
    ('left ear', 'left eye'),
    ('right ear', 'right eye'),
    ('left eye', 'right eye'),
    ('left shoulder', 'right shoulder'),
    ('left shoulder', 'left elbow'),
    ('left shoulder', 'left hip'),
    ('right shoulder', 'right elbow'),
    ('right shoulder', 'right hip'),
    ('left elbow', 'left wrist'),
    ('right elbow', 'right wrist'),
    ('left hip', 'right hip'),
    ('left hip', 'left knee'),
    ('right hip', 'right knee'),
    ('left knee', 'left ankle'),
    ('right knee', 'right ankle'),
)

def shadow_text(dwg, x, y, text, font_size=16):
    # draw text, half opacity
    dwg.text((10,10), text, fill=(255,255,255,128))

def draw_pose(dwg, pose, src_size, color='yellow', threshold=0.2):
    box_x, box_y, box_w, box_h = (0, 0, src_size[0], src_size[1])
    scale_x, scale_y = src_size[0] / box_w, src_size[1] / box_h
    xys = {}
    for label, keypoint in pose.keypoints.items():
        if keypoint.score < threshold: continue
        # Offset and scale to source coordinate space.
        kp_y = int((keypoint.yx[0] - box_y) * scale_y)
        kp_x = int((keypoint.yx[1] - box_x) * scale_x)

        xys[label] = (kp_x, kp_y)
        dwg.ellipse([(kp_x, kp_y), (kp_x + 5, kp_y + 5)], fill='cyan')

    for a, b in EDGES:
        if a not in xys or b not in xys: continue
        ax, ay = xys[a]
        bx, by = xys[b]
        dwg.line([(ax, ay),(bx, by)], fill=color, width=2)

def render_overlay(pil_image, poses, src_size, inference_time):
        text_line = 'PoseNet: %.1fms Nposes %d' % (
            inference_time, len(poses))
        shadow_text(pil_image, 10, 20, text_line)
        for pose in poses:
            draw_pose(pil_image, pose, src_size)

engine = PoseEngine('models/posenet_mobilenet_v1_075_481_641_quant_decoder_edgetpu.tflite')

def process_poses(pil_image):
    pil_image.resize((641, 481), Image.NEAREST)
    poses, inference_time = engine.DetectPosesInImage(np.uint8(pil_image))
    print('Inference time: %.fms' % inference_time)

    draw = ImageDraw.Draw(pil_image)
    render_overlay(draw, poses, (641, 481), inference_time)
    del draw

    for pose in poses:
        if pose.score < 0.4: continue
        print('\nPose Score: ', pose.score)
        for label, keypoint in pose.keypoints.items():
            print(' %-20s x=%-4d y=%-4d score=%.1f' %
                  (label, keypoint.yx[1], keypoint.yx[0], keypoint.score))
    return poses

if __name__ == "__main__":
    os.system('wget https://upload.wikimedia.org/wikipedia/commons/thumb/3/38/'
              'Hindu_marriage_ceremony_offering.jpg/'
              '640px-Hindu_marriage_ceremony_offering.jpg -O couple.jpg')
    pil_image = Image.open('couple.jpg')
    process_poses(pil_image)
    pil_image.save('img.jpg')


text:
from sspipe import p, px
import numpy as np

def test_scalar_rhs():
    assert np.int32(1) | p(lambda x: x + 1) | (px == 2)

def test_scalar_lhs():
    assert 2 | px + np.int32(1)

def test_rhs():
    assert np.array([1, 2]) | p(lambda x: x.sum()) | (px == 3)

def test_rhs_px():
    assert np.array([1, 2]) | (px.sum() == 3)

def test_lhs():
    assert 2 | p(np.log2) | (px == 1)

def test_lhs_px():
    assert 2 | np.power(px, px + 1) | (px == 8)

max_stars_repo_name: python-pipe/hellp

max_issues_repo_name: python-pipe/hellp

max_forks_repo_name: python-pipe/hellp