text:
import time
import json
import datetime as dt
import numpy as np
import pandas as pd
import pytest

from solarforecastarbiter import datamodel
from solarforecastarbiter.metrics import preprocessing

THREE_HOURS = pd.date_range(start='2019-03-31T12:00:00',
                            periods=3,
                            freq='60min',
                            tz='MST',
                            name='timestamp')

THREE_HOUR_SERIES = pd.Series(np.arange(1., 4., 1.), index=THREE_HOURS,
                              name='value')

THREE_HOUR_NAN_SERIES = pd.Series([1.0, np.nan, 4.0], index=THREE_HOURS,
                                  name="value")

THREE_HOURS_NAN = THREE_HOURS[[True, False, True]]

THREE_HOURS_EMPTY = pd.DatetimeIndex([], name="timestamp", freq="60min",
                                     tz="MST")

THREE_HOUR_EMPTY_SERIES = pd.Series([], index=THREE_HOURS_EMPTY, name="value",
                                    dtype='float64')

EMPTY_OBJ_SERIES = pd.Series(
    [],
    dtype=object,
    name="value",
    index=pd.DatetimeIndex([], freq="10min", tz="MST", name="timestamp")
)

THIRTEEN_10MIN = pd.date_range(start='2019-03-31T12:00:00',
                               periods=13,
                               freq='10min',
                               tz='MST',
                               name='timestamp')

THIRTEEN_10MIN_SERIES = pd.Series((np.arange(0., 13., 1.)/6)+1,
                                  index=THIRTEEN_10MIN)

# Bitwise-flag integers (only test validated and versioned data)
NT_UF = int(0b10011)  # Nighttime, User Flagged and version 0 (19)
CSE_NT = int(0b1000010010)  # Clearsky exceeded, nighttime, and version 0 (530)
CSE = int(0b1000000010)  # Clearsky exceeded and version 0 (514)
OK = int(0b10)  # OK version 0 (2)

def create_preprocessing_result(counts):
    """Create preprocessing results in order as
    1. Forecast Undefined Values,
    2. Forecast Values Discarded by Alignment,
    3. Observation Undefined Values,
    4. Observation Values Discarded by Alignment
    """
    return {
        "Forecast " + preprocessing.UNDEFINED_DATA_STRING: counts[0],
        "Forecast " + preprocessing.DISCARD_DATA_STRING: counts[1],
        "Observation " + preprocessing.UNDEFINED_DATA_STRING: counts[2],
        "Observation " + preprocessing.DISCARD_DATA_STRING: counts[3]
    }

@pytest.mark.parametrize('obs_interval_label',
                         ['beginning', 'instant', 'ending'])
@pytest.mark.parametrize('fx_interval_label',
                         ['beginning', 'ending'])
@pytest.mark.parametrize('fx_series,obs_series,expected_dt,expected_res', [
    (THREE_HOUR_SERIES, THREE_HOUR_SERIES, THREE_HOURS, [0]*4),
    (THREE_HOUR_SERIES, THIRTEEN_10MIN_SERIES, THREE_HOURS, [0]*4),
    (THIRTEEN_10MIN_SERIES, THIRTEEN_10MIN_SERIES, THIRTEEN_10MIN, [0]*4),
    (THREE_HOUR_SERIES, THREE_HOUR_NAN_SERIES, THREE_HOURS_NAN, [0, 1, 1, 0]),
    (THREE_HOUR_NAN_SERIES, THREE_HOUR_SERIES, THREE_HOURS_NAN, [1, 0, 0, 1]),
    (THREE_HOUR_NAN_SERIES, THREE_HOUR_NAN_SERIES, THREE_HOURS_NAN, [1, 0, 1, 0]),  # NOQA
    (THREE_HOUR_SERIES, THREE_HOUR_EMPTY_SERIES, THREE_HOURS_EMPTY, [0, 3, 0, 0]),  # NOQA
    (THREE_HOUR_EMPTY_SERIES, THREE_HOUR_SERIES, THREE_HOURS_EMPTY, [0, 0, 0, 3]),  # NOQA
    (THREE_HOUR_SERIES, EMPTY_OBJ_SERIES, THREE_HOURS_EMPTY, [0, 3, 0, 0]),
])
def test_resample_and_align(
        site_metadata, obs_interval_label, fx_interval_label, fx_series,
        obs_series, expected_dt, expected_res):
    # Create the ForecastObservation to match interval_lengths of data
    observation = datamodel.Observation(
        site=site_metadata, name='dummy obs', variable='ghi',
        interval_value_type='instantaneous', uncertainty=1,
        interval_length=pd.Timedelta(obs_series.index.freq),
        interval_label=obs_interval_label
    )
    forecast = datamodel.Forecast(
        site=site_metadata, name='dummy fx', variable='ghi',
        interval_value_type='instantaneous',
        interval_length=pd.Timedelta(fx_series.index.freq),
        interval_label=fx_interval_label,
        issue_time_of_day=dt.time(hour=5),
        lead_time_to_start=pd.Timedelta('1h'),
        run_length=pd.Timedelta('12h')
    )
    fx_obs = datamodel.ForecastObservation(forecast=forecast,
                                           observation=observation)

    # Use local tz
    local_tz = f"Etc/GMT{int(time.timezone/3600):+d}"

    fx_vals, obs_vals, ref_vals, res_dict = preprocessing.resample_and_align(
            fx_obs, fx_series, obs_series, None, local_tz)

    # Localize datetimeindex
    expected_dt = expected_dt.tz_convert(local_tz)

    pd.testing.assert_index_equal(fx_vals.index,
                                  obs_vals.index,
                                  check_categorical=False)
    pd.testing.assert_index_equal(obs_vals.index,
                                  expected_dt,
                                  check_categorical=False)

    expected_result = create_preprocessing_result(expected_res)
    assert res_dict == expected_result

def test_resample_and_align_fx_aggregate(single_forecast_aggregate):
    fx_series = THREE_HOUR_SERIES
    obs_series = THREE_HOUR_SERIES
    fx_values, obs_values, ref, res_dict = preprocessing.resample_and_align(
        single_forecast_aggregate, fx_series, obs_series, None, 'UTC')

    pd.testing.assert_index_equal(fx_values.index,
                                  obs_values.index,
                                  check_categorical=False)
    pd.testing.assert_index_equal(obs_values.index,
                                  THREE_HOURS,
                                  check_categorical=False)

@pytest.mark.parametrize("label_obs,label_fx,length_obs,length_fx,freq", [
    ("beginning", "ending", "5min", "5min", "5min"),
    ("beginning", "ending", "5min", "1h", "1h"),
    ("ending", "beginning", "5min", "5min", "5min"),
    ("ending", "beginning", "5min", "1h", "1h"),
])
def test_resample_and_align_interval_label(site_metadata, label_obs, label_fx,
                                           length_obs, length_fx, freq):

    observation = datamodel.Observation(
        site=site_metadata, name='dummy obs', variable='ghi',
        interval_value_type='instantaneous', uncertainty=1,
        interval_length=pd.Timedelta(length_obs),
        interval_label=label_obs
    )
    forecast = datamodel.Forecast(
        site=site_metadata, name='dummy fx', variable='ghi',
        interval_value_type='instantaneous',
        interval_length=pd.Timedelta(length_fx),
        interval_label=label_fx,
        issue_time_of_day=dt.time(hour=5),
        lead_time_to_start=pd.Timedelta('1h'),
        run_length=pd.Timedelta('12h')
    )
    fx_obs = datamodel.ForecastObservation(forecast=forecast,
                                           observation=observation)

    ts_obs = pd.date_range(
        start='2019-03-31T12:00:00',
        end='2019-05-01T00:00:00',
        freq=length_obs,
        tz='MST',
        name='timestamp'
    )
    ts_fx = pd.date_range(
        start='2019-03-31T12:00:00',
        end='2019-05-01T00:00:00',
        freq=length_fx,
        tz='MST',
        name='timestamp'
    )
    obs_series = pd.Series(index=ts_obs, data=np.random.rand(len(ts_obs)) + 10)
    fx_series = pd.Series(index=ts_fx, data=np.random.rand(len(ts_fx)) + 10)

    # Use local tz
    local_tz = f"Etc/GMT{int(time.timezone/3600):+d}"

    fx_out, obs_out, ref_out, res_dict = preprocessing.resample_and_align(
        fx_obs, fx_series, obs_series, None, local_tz)

    assert obs_out.index.freq == freq

@pytest.mark.parametrize("interval_label", ["beginning", "instant", "ending"])
@pytest.mark.parametrize("tz,local_tz,local_ts", [
    ("UTC", "UTC", ["20190702T0000", "20190702T0100"]),
    ("UTC", "US/Pacific", ["20190701T1700", "20190701T1800"]),
    ("US/Pacific", "UTC", ["20190702T0700", "20190702T0800"]),
    ("US/Central", "US/Pacific", ["20190701T2200", "20190701T2300"]),
    ("US/Pacific", "US/Eastern", ["20190702T0300", "20190702T0400"]),
])
def test_resample_and_align_timezone(site_metadata, interval_label, tz,
                                     local_tz, local_ts):

    expected_dt = pd.DatetimeIndex(local_ts, tz=local_tz)

    # Create the fx/obs pair
    ts = pd.DatetimeIndex(["20190702T0000", "20190702T0100"], tz=tz)
    fx_series = pd.Series([1.0, 4.0], index=ts, name="value")
    obs_series = pd.Series([1.1, 2.7], index=ts, name="value")
    observation = datamodel.Observation(
        site=site_metadata, name='dummy obs', variable='ghi',
        interval_value_type='instantaneous', uncertainty=1,
        interval_length=pd.Timedelta(obs_series.index.freq),
        interval_label=interval_label,
    )
    forecast = datamodel.Forecast(
        site=site_metadata, name='dummy fx', variable='ghi',
        interval_value_type='instantaneous',
        interval_length=pd.Timedelta(fx_series.index.freq),
        interval_label=interval_label,
        issue_time_of_day=dt.time(hour=5),
        lead_time_to_start=pd.Timedelta('1h'),
        run_length=pd.Timedelta('12h')
    )
    fx_obs = datamodel.ForecastObservation(forecast=forecast,
                                           observation=observation)

    fx_values, obs_values, ref_values, _ = preprocessing.resample_and_align(
        fx_obs, fx_series, obs_series, None, local_tz)

    pd.testing.assert_index_equal(fx_values.index,
                                  obs_values.index,
                                  check_categorical=False)
    pd.testing.assert_index_equal(obs_values.index,
                                  expected_dt,
                                  check_categorical=False)

def test_resample_and_align_with_ref(
        single_forecast_observation_reffx):
    tz = 'UTC'
    fx_obs = single_forecast_observation_reffx
    fx_series = THREE_HOUR_SERIES
    ref_series = THREE_HOUR_SERIES
    obs_series = THREE_HOUR_SERIES
    fx_values, obs_values, ref_values, _ = preprocessing.resample_and_align(
        fx_obs, fx_series, obs_series, ref_series, tz)
    pd.testing.assert_index_equal(fx_values.index,
                                  obs_values.index,
                                  check_categorical=False)
    pd.testing.assert_index_equal(fx_values.index,
                                  ref_values.index,
                                  check_categorical=False)
    pd.testing.assert_index_equal(obs_values.index,
                                  THREE_HOURS,
                                  check_categorical=False)

def test_resample_and_align_ref_less_fx(
        single_forecast_observation_reffx):
    tz = 'UTC'
    fx_obs = single_forecast_observation_reffx
    nine_hour_series = pd.concat([
        THREE_HOUR_SERIES.shift(periods=3, freq='1h'),
        THREE_HOUR_SERIES,
        THREE_HOUR_SERIES.shift(periods=3, freq='-1h')])
    fx_series = THREE_HOUR_SERIES
    ref_series = nine_hour_series
    obs_series = nine_hour_series
    fx_values, obs_values, ref_values, _ = preprocessing.resample_and_align(
        fx_obs, fx_series, obs_series, ref_series, tz)
    pd.testing.assert_index_equal(fx_values.index,
                                  obs_values.index,
                                  check_categorical=False)
    pd.testing.assert_index_equal(fx_values.index,
                                  ref_values.index,
                                  check_categorical=False)
    pd.testing.assert_index_equal(obs_values.index,
                                  THREE_HOURS,
                                  check_categorical=False)

def test_resample_and_align_ref_error_None(
        single_forecast_observation, single_forecast_observation_reffx):
    tz = 'UTC'

    # no ref_fx object, but supplied ref_fx series
    fx_obs = single_forecast_observation
    fx_series = THREE_HOUR_SERIES
    ref_series = THREE_HOUR_SERIES
    obs_series = THREE_HOUR_SERIES
    with pytest.raises(ValueError):
        preprocessing.resample_and_align(
            fx_obs, fx_series, obs_series, ref_series, tz)

    # ref_fx object, but no supplied ref_fx series
    fx_obs = single_forecast_observation_reffx
    fx_series = THREE_HOUR_SERIES
    ref_series = None
    obs_series = THREE_HOUR_SERIES
    with pytest.raises(ValueError):
        preprocessing.resample_and_align(
            fx_obs, fx_series, obs_series, ref_series, tz)

@pytest.mark.parametrize('attr,value', [
    ('interval_label', 'ending'),
    ('interval_length', pd.Timedelta('20min')),
])
def test_resample_and_align_ref_error(
        single_forecast_observation_reffx, attr, value):
    tz = 'UTC'

    changes = {attr: value}
    # ref_fx object parameters are inconsistent with fx object parameters
    ref_fx = single_forecast_observation_reffx.reference_forecast.replace(
        **changes)
    fx_obs = single_forecast_observation_reffx.replace(
        reference_forecast=ref_fx)
    fx_series = THREE_HOUR_SERIES
    ref_series = THREE_HOUR_SERIES
    obs_series = THREE_HOUR_SERIES
    with pytest.raises(ValueError):
        preprocessing.resample_and_align(
            fx_obs, fx_series, obs_series, ref_series, tz)

def test_resample_and_align_ref_error_prob(prob_forecasts, single_observation):
    tz = 'UTC'
    cv = prob_forecasts.constant_values[0].replace(axis='y')
    ref_fx = prob_forecasts.replace(axis='y', constant_values=(cv,))
    fx_obs = datamodel.ForecastObservation(
        prob_forecasts,
        single_observation,
        reference_forecast=ref_fx)
    fx_series = THREE_HOUR_SERIES
    ref_series = THREE_HOUR_SERIES
    obs_series = THREE_HOUR_SERIES
    with pytest.raises(ValueError):
        preprocessing.resample_and_align(
            fx_obs, fx_series, obs_series, ref_series, tz)

def test_resample_and_align_prob(prob_forecasts, single_observation):
    tz = 'UTC'
    fx_obs = datamodel.ForecastObservation(
        prob_forecasts,
        single_observation,
        reference_forecast=prob_forecasts.replace(
            name='reference'))
    fx_data = THREE_HOUR_SERIES.to_frame()
    obs_data = THREE_HOUR_SERIES
    ref_data = THREE_HOUR_SERIES.to_frame()
    fx_values, obs_values, ref_values, _ = preprocessing.resample_and_align(
        fx_obs, fx_data, obs_data, ref_data, tz)
    pd.testing.assert_index_equal(fx_values.index,
                                  ref_values.index,
                                  check_categorical=False)
    pd.testing.assert_index_equal(fx_values.index,
                                  obs_values.index,
                                  check_categorical=False)
    pd.testing.assert_index_equal(obs_values.index,
                                  THREE_HOURS,
                                  check_categorical=False)

def test_resample_and_align_prob_constant_value(
        prob_forecast_constant_value, single_observation):
    tz = 'UTC'
    fx_obs = datamodel.ForecastObservation(
        prob_forecast_constant_value,
        single_observation,
        reference_forecast=prob_forecast_constant_value.replace(
            name='reference'))
    fx_data = THREE_HOUR_SERIES
    obs_data = THREE_HOUR_SERIES
    ref_data = THREE_HOUR_SERIES
    fx_values, obs_values, ref_values, _ = preprocessing.resample_and_align(
        fx_obs, fx_data, obs_data, ref_data, tz)
    pd.testing.assert_index_equal(fx_values.index,
                                  ref_values.index,
                                  check_categorical=False)
    pd.testing.assert_index_equal(fx_values.index,
                                  obs_values.index,
                                  check_categorical=False)
    pd.testing.assert_index_equal(obs_values.index,
                                  THREE_HOURS,
                                  check_categorical=False)

@pytest.mark.parametrize('obs,somecounts', [
    (pd.DataFrame(index=pd.DatetimeIndex([], name='timestamp'),
                  columns=['value', 'quality_flag']),
     {'USER FLAGGED': 0}),
    (pd.DataFrame({'value': [1., 2., 3.],
                   'quality_flag': [OK, OK, OK]},
                  index=THREE_HOURS),
     {'USER FLAGGED': 0}),
    (pd.DataFrame({'value': [1., 2., 3.],
                   'quality_flag': [NT_UF, NT_UF, NT_UF]},
                  index=THREE_HOURS),
     {'NIGHTTIME': 3, 'USER FLAGGED': 3, 'STALE VALUES': 0}),
    (pd.DataFrame({'value': [1., 2., 3.],
                   'quality_flag': [CSE_NT, CSE, OK]},
                  index=THREE_HOURS),
     {'NIGHTTIME': 1, 'USER FLAGGED': 0, 'CLEARSKY EXCEEDED': 2}),
])
@pytest.mark.parametrize('handle_func', [preprocessing.exclude])
@pytest.mark.parametrize('filter_', [
    datamodel.QualityFlagFilter(
        (
            "USER FLAGGED",
            "NIGHTTIME",
            "LIMITS EXCEEDED",
            "STALE VALUES",
            "INTERPOLATED VALUES",
            "INCONSISTENT IRRADIANCE COMPONENTS",
        )
    ),
    pytest.param(
        datamodel.TimeOfDayFilter((dt.time(12, 0),
                                   dt.time(14, 0))),
        marks=pytest.mark.xfail(strict=True, type=TypeError)
    )
])
def test_apply_validation(obs, handle_func, filter_, somecounts):
    result, counts = preprocessing.apply_validation(obs, filter_,
                                                    handle_func)

    # Check length and timestamps of observation
    assert len(obs[obs.quality_flag.isin([OK, CSE])]) == \
        len(result)
    if not result.empty:
        assert obs[obs.quality_flag.isin([OK, CSE])].index.equals(
            result.index)
    assert set(filter_.quality_flags) == set(counts.keys())
    for k, v in somecounts.items():
        assert counts.get(k, v) == v

@pytest.mark.parametrize('values,qflags,expectation', [
    (THREE_HOUR_SERIES, None, THREE_HOUR_SERIES),
    (THREE_HOUR_SERIES,
        pd.DataFrame({1: [0, 0, 0]}, index=THREE_HOURS),
        THREE_HOUR_SERIES),
    (THREE_HOUR_SERIES,
        pd.DataFrame({1: [0, 1, 0]}, index=THREE_HOURS),
        THREE_HOUR_SERIES[[True, False, True]]),
    (THREE_HOUR_SERIES,
        pd.DataFrame({1: [1, 1, 0], 2: [0, 1, 0]}, index=THREE_HOURS),
        THREE_HOUR_SERIES[[False, False, True]]),
    (THREE_HOUR_SERIES,
        pd.DataFrame({1: [1, 1, 1], 2: [1, 1, 1]}, index=THREE_HOURS),
        THREE_HOUR_SERIES[[False, False, False]]),
    (pd.Series([np.NaN, np.NaN, np.NaN], index=THREE_HOURS),
        None,
        THREE_HOUR_SERIES[[False, False, False]]),
    (pd.Series([1., np.NaN, 3.], index=THREE_HOURS),
        pd.DataFrame({1: [0, 0, 0], 2: [0, 0, 0]}, index=THREE_HOURS),
        THREE_HOUR_SERIES[[True, False, True]]),
    (pd.Series([1., np.NaN, 3.], index=THREE_HOURS),
        pd.DataFrame({1: [1, 1, 0], 2: [1, 0, 0]}, index=THREE_HOURS),
        THREE_HOUR_SERIES[[False, False, True]])
])
def test_exclude(values, qflags, expectation):
    result = preprocessing.exclude(values, qflags)
    pd.testing.assert_series_equal(result, expectation, check_names=False)

def test_merge_quality_filters():
    filters = [
        datamodel.QualityFlagFilter(('USER FLAGGED', 'NIGHTTIME',
                                     'CLIPPED VALUES')),
        datamodel.QualityFlagFilter(('SHADED', 'NIGHTTIME',)),
        datamodel.QualityFlagFilter(())
    ]
    out = preprocessing._merge_quality_filters(filters)
    assert set(out.quality_flags) == {'USER FLAGGED', 'NIGHTTIME',
                                      'CLIPPED VALUES', 'SHADED'}

def test_process_forecast_observations(report_objects, quality_filter,
                                       timeofdayfilter, mocker):
    report, observation, forecast_0, forecast_1, aggregate, forecast_agg = report_objects  # NOQA
    forecast_ref = report.report_parameters.object_pairs[1].reference_forecast
    obs_ser = pd.Series(np.arange(8),
                        index=pd.date_range(start='2019-03-31T12:00:00',
                                            periods=8,
                                            freq='15min',
                                            tz='MST',
                                            name='timestamp'))
    obs_df = obs_ser.to_frame('value')
    obs_df['quality_flag'] = OK
    agg_df = THREE_HOUR_SERIES.to_frame('value')
    agg_df['quality_flag'] = OK
    data = {
        observation: obs_df,
        forecast_0: THREE_HOUR_SERIES,
        forecast_1: THREE_HOUR_SERIES,
        forecast_ref: THREE_HOUR_SERIES,
        forecast_agg: THREE_HOUR_SERIES,
        aggregate: agg_df
    }
    filters = [quality_filter, timeofdayfilter]
    logger = mocker.patch('solarforecastarbiter.metrics.preprocessing.logger')
    processed_fxobs_list = preprocessing.process_forecast_observations(
        report.report_parameters.object_pairs, filters, data, 'MST')
    assert len(processed_fxobs_list) == len(
        report.report_parameters.object_pairs)
    assert logger.warning.called
    assert not logger.error.called
    for proc_fxobs in processed_fxobs_list:
        assert isinstance(proc_fxobs, datamodel.ProcessedForecastObservation)
        assert all(isinstance(vr, datamodel.ValidationResult)
                   for vr in proc_fxobs.validation_results)
        assert all(isinstance(pr, datamodel.PreprocessingResult)
                   for pr in proc_fxobs.preprocessing_results)
        assert isinstance(proc_fxobs.forecast_values, pd.Series)
        assert isinstance(proc_fxobs.observation_values, pd.Series)
        pd.testing.assert_index_equal(proc_fxobs.forecast_values.index,
                                      proc_fxobs.observation_values.index)

def test_process_probabilistic_forecast_observations(
        cdf_and_cv_report_objects, quality_filter, timeofdayfilter, mocker):
    report, observation, cdf_forecast_0, cdf_forecast_1, \
        cv_forecast_0, cv_forecast_1, aggregate, cdf_forecast_agg, \
        cv_forecast_agg = cdf_and_cv_report_objects
    cdf_forecast_ref = report.report_parameters.object_pairs[1].reference_forecast  # NOQA
    cv_forecast_ref = report.report_parameters.object_pairs[3].reference_forecast  # NOQA
    periods_1min = 60*24*7
    obs_ser = pd.Series(np.arange(periods_1min)/60.,
                        index=pd.date_range(start='2020-04-01T00:00:00',
                                            periods=periods_1min,
                                            freq='1min',
                                            tz='MST',
                                            name='timestamp'))
    cdf_fx_df = pd.DataFrame({'25': np.arange(periods_1min/60),
                              '50': np.arange(periods_1min/60)+1,
                              '75': np.arange(periods_1min/60)+2},
                             index=pd.date_range(start='2020-04-01T00:00:00',
                                                 periods=periods_1min/60,
                                                 freq='60min',
                                                 tz='MST',
                                                 name='timestamp'))
    obs_df = obs_ser.to_frame('value')
    obs_df['quality_flag'] = OK
    agg_df = obs_df.resample('60T').asfreq()
    agg_df['quality_flag'] = OK

    data = {
        observation: obs_df,
        aggregate: agg_df,
        cdf_forecast_0: cdf_fx_df,
        cdf_forecast_1: cdf_fx_df,
        cdf_forecast_agg: cdf_fx_df,
        cdf_forecast_ref: cdf_fx_df,
        cv_forecast_0: cdf_fx_df['25'],
        cv_forecast_1: cdf_fx_df['50'],
        cv_forecast_agg: cdf_fx_df['75'],
        cv_forecast_ref: cdf_fx_df['50']
    }

    filters = [quality_filter, timeofdayfilter]
    logger = mocker.patch('solarforecastarbiter.metrics.preprocessing.logger')
    processed_fxobs_list = preprocessing.process_forecast_observations(
        report.report_parameters.object_pairs, filters, data, 'MST')
    assert len(processed_fxobs_list) == len(
        report.report_parameters.object_pairs)
    assert logger.warning.called
    assert not logger.error.called
    for proc_fxobs in processed_fxobs_list:
        assert isinstance(proc_fxobs, datamodel.ProcessedForecastObservation)
        assert all(isinstance(vr, datamodel.ValidationResult)
                   for vr in proc_fxobs.validation_results)
        assert all(isinstance(pr, datamodel.PreprocessingResult)
                   for pr in proc_fxobs.preprocessing_results)
        if isinstance(proc_fxobs.original.forecast,
                      datamodel.ProbabilisticForecast):
            assert isinstance(proc_fxobs.forecast_values, pd.DataFrame)
        else:
            assert isinstance(proc_fxobs.forecast_values, pd.Series)
        assert isinstance(proc_fxobs.observation_values, pd.Series)
        pd.testing.assert_index_equal(proc_fxobs.forecast_values.index,
                                      proc_fxobs.observation_values.index)

def test_process_forecast_observations_no_data(
        report_objects, quality_filter, mocker):
    report, observation, forecast_0, forecast_1, aggregate, forecast_agg = report_objects  # NOQA
    forecast_ref = report.report_parameters.object_pairs[1].reference_forecast
    agg_df = THREE_HOUR_SERIES.to_frame('value')
    agg_df['quality_flag'] = NT_UF
    data = {
        forecast_0: THREE_HOUR_SERIES,
        forecast_1: THREE_HOUR_SERIES,
        forecast_ref: THREE_HOUR_SERIES,
        forecast_agg: THREE_HOUR_SERIES,
        aggregate: agg_df
    }
    filters = [quality_filter]
    logger = mocker.patch('solarforecastarbiter.metrics.preprocessing.logger')
    processed_fxobs_list = preprocessing.process_forecast_observations(
        report.report_parameters.object_pairs, filters, data, 'MST')
    assert len(processed_fxobs_list) == len(
        report.report_parameters.object_pairs)
    assert logger.error.called
    for proc_fxobs in processed_fxobs_list:
        assert isinstance(proc_fxobs, datamodel.ProcessedForecastObservation)
        assert isinstance(proc_fxobs.forecast_values, pd.Series)
        assert isinstance(proc_fxobs.observation_values, pd.Series)
        pd.testing.assert_index_equal(proc_fxobs.forecast_values.index,
                                      proc_fxobs.observation_values.index)

def test_process_forecast_observations_resample_fail(
        report_objects, quality_filter, mocker):
    report, observation, forecast_0, forecast_1, aggregate, forecast_agg = report_objects  # NOQA
    forecast_ref = report.report_parameters.object_pairs[1].reference_forecast
    obs_ser = pd.Series(np.arange(8),
                        index=pd.date_range(start='2019-03-31T12:00:00',
                                            periods=8,
                                            freq='15min',
                                            tz='MST',
                                            name='timestamp'))
    obs_df = obs_ser.to_frame('value')
    obs_df['quality_flag'] = OK
    agg_df = THREE_HOUR_SERIES.to_frame('value')
    agg_df['quality_flag'] = OK
    data = {
        observation: obs_df,
        forecast_0: THREE_HOUR_SERIES,
        forecast_1: THREE_HOUR_SERIES,
        forecast_ref: THREE_HOUR_SERIES,
        forecast_agg: THREE_HOUR_SERIES,
        aggregate: agg_df
    }
    filters = [quality_filter]
    logger = mocker.patch('solarforecastarbiter.metrics.preprocessing.logger')
    mocker.patch(
        'solarforecastarbiter.metrics.preprocessing.resample_and_align',
        side_effect=ValueError)
    processed_fxobs_list = preprocessing.process_forecast_observations(
        report.report_parameters.object_pairs, filters, data, 'MST')
    assert len(processed_fxobs_list) == 0
    assert logger.error.called

def test_process_forecast_observations_same_name(
        report_objects, quality_filter, mocker):
    report, observation, forecast_0, forecast_1, aggregate, forecast_agg = report_objects  # NOQA
    forecast_ref = report.report_parameters.object_pairs[1].reference_forecast
    obs_ser = pd.Series(np.arange(8),
                        index=pd.date_range(start='2019-03-31T12:00:00',
                                            periods=8,
                                            freq='15min',
                                            tz='MST',
                                            name='timestamp'))
    obs_df = obs_ser.to_frame('value')
    obs_df['quality_flag'] = OK
    agg_df = THREE_HOUR_SERIES.to_frame('value')
    agg_df['quality_flag'] = OK
    data = {
        observation: obs_df,
        forecast_0: THREE_HOUR_SERIES,
        forecast_1: THREE_HOUR_SERIES,
        forecast_ref: THREE_HOUR_SERIES,
        forecast_agg: THREE_HOUR_SERIES,
        aggregate: agg_df
    }
    filters = [quality_filter]
    fxobs = report.report_parameters.object_pairs
    fxobs = list(fxobs) + [fxobs[0], fxobs[0]]
    processed_fxobs_list = preprocessing.process_forecast_observations(
        fxobs, filters, data, 'MST')
    assert len(processed_fxobs_list) == len(
        fxobs)
    assert len(set(pfxobs.name for pfxobs in processed_fxobs_list)) == len(
        fxobs)

def test_name_pfxobs_recursion_limit():
    name = 'whoami'
    cn = [name]
    cn += [f'{name}-{i:02d}' for i in range(129)]
    assert preprocessing._name_pfxobs(cn, name) == f'{name}-99'

@pytest.mark.parametrize("obs_index,fx_index,expected_dt", [
    (THREE_HOURS, THREE_HOURS, THREE_HOURS),
    (THIRTEEN_10MIN, THIRTEEN_10MIN, THIRTEEN_10MIN),

    # mismatch interval_length
    pytest.param(THREE_HOURS, THIRTEEN_10MIN, [],
                 marks=pytest.mark.xfail(strict=True, type=ValueError)),
    pytest.param(THIRTEEN_10MIN, THREE_HOURS, [],
                 marks=pytest.mark.xfail(strict=True, type=ValueError)),
])
def test_resample_and_align_event(single_event_forecast_observation,
                                  obs_index, fx_index, expected_dt):

    obs_values = np.random.randint(0, 2, size=len(obs_index), dtype=bool)
    fx_values = np.random.randint(0, 2, size=len(fx_index), dtype=bool)
    obs_series = pd.Series(obs_values, index=obs_index)
    fx_series = pd.Series(fx_values, index=fx_index)

    fxobs = single_event_forecast_observation

    # Use local tz
    local_tz = f"Etc/GMT{int(time.timezone/3600):+d}"

    fx_vals, obs_vals, ref_vals, results = preprocessing.resample_and_align(
        fxobs, fx_series, obs_series, None, local_tz
    )

    assert isinstance(results, dict)

    # Localize datetimeindex
    expected_dt = expected_dt.tz_convert(local_tz)

    pd.testing.assert_index_equal(fx_vals.index,
                                  obs_vals.index,
                                  check_categorical=False)
    pd.testing.assert_index_equal(obs_vals.index,
                                  expected_dt,
                                  check_categorical=False)

@pytest.mark.parametrize("fx_interval_length, obs_interval_length", [
    (60, 60),
    (5, 5),
    pytest.param(30, 10,
                 marks=pytest.mark.xfail(strict=True, type=ValueError)),
    pytest.param(10, 15,
                 marks=pytest.mark.xfail(strict=True, type=ValueError)),
])
def test__resample_event_obs(single_site, single_event_forecast_text,
                             single_event_observation_text,
                             fx_interval_length, obs_interval_length):

    fx_dict = json.loads(single_event_forecast_text)
    fx_dict['site'] = single_site
    fx_dict.update({"interval_length": fx_interval_length})
    fx = datamodel.EventForecast.from_dict(fx_dict)

    obs_dict = json.loads(single_event_observation_text)
    obs_dict['site'] = single_site
    obs_dict.update({"interval_length": obs_interval_length})
    obs = datamodel.Observation.from_dict(obs_dict)

    freq = pd.Timedelta(f"{obs_interval_length}min")
    index = pd.date_range(start="20200301T00Z", end="20200304T00Z", freq=freq)
    obs_series = pd.Series(np.random.randint(0, 2, len(index), dtype=bool),
                           index=index)

    obs_resampled = preprocessing._resample_event_obs(fx, obs, obs_series)
    pd.testing.assert_index_equal(obs_series.index, obs_resampled.index,
                                  check_categorical=False)

@pytest.mark.parametrize("data", [
    [True, True, False],
    [1, 1, 0],
    [1.0, 1.0, 0.0],
    pytest.param([0, 1, 2],
                 marks=pytest.mark.xfail(strict=True, type=TypeError)),
    pytest.param([0.0, 1.0, 2.0],
                 marks=pytest.mark.xfail(strict=True, type=TypeError)),
])
def test__validate_event_dtype(data):
    ser = pd.Series(data)
    print(ser.head())
    ser_conv = preprocessing._validate_event_dtype(ser)

    pd.testing.assert_index_equal(ser.index, ser_conv.index,
                                  check_categorical=False)
    assert ser_conv.dtype == bool

max_stars_repo_name: awig/solarforecastarbiter-core

max_issues_repo_name: awig/solarforecastarbiter-core

max_forks_repo_name: awig/solarforecastarbiter-core


text:
#!/usr/bin/env python
# vim: set fileencoding=utf-8 ts=4 sts=4 sw=4 et tw=80 :
#
# Compute order spacing with and without accounting for half the cross-
# dispersion occurring prior to grating. Effectively, the "accounting" case
# allows the gamma angle to vary with wavelength whereas the naive case uses
# a fixed gamma angle for all wavelengths.
#
# Rob Siverd
# Created:       2019-01-20
# Last modified: 2019-02-08
#--------------------------------------------------------------------------
#**************************************************************************
#--------------------------------------------------------------------------

## Current version:
__version__ = "0.1.5"

## Optional matplotlib control:
#from matplotlib import use, rc, rcParams
#from matplotlib import use
#from matplotlib import rc
#from matplotlib import rcParams
#use('GTKAgg')  # use GTK with Anti-Grain Geometry engine
#use('agg')     # use Anti-Grain Geometry engine (file only)
#use('ps')      # use PostScript engine for graphics (file only)
#use('cairo')   # use Cairo (pretty, file only)
#rc('font',**{'family':'sans-serif','sans-serif':['Helvetica']})
#rc('font',**{'family':'serif','serif':['Palatino']})
#rc('font',**{'sans-serif':'Arial','family':'sans-serif'})
#rc('text', usetex=True) # enables text rendering with LaTeX (slow!)
#rcParams['axes.formatter.useoffset'] = False   # v. 1.4 and later

## Python version-agnostic module reloading:
try:
    reload                              # Python 2.7
except NameError:
    try:
        from importlib import reload    # Python 3.4+
    except ImportError:
        from imp import reload          # Python 3.0 - 3.3

## Modules:
#import argparse
#import mimetypes
#import linecache
#import getopt
#import shutil
#import resource
#import signal
#import glob
#import gc
import os
import sys
import time
#import vaex
#import calendar
#import ephem
import numpy as np
#from numpy.lib.recfunctions import append_fields
#import datetime as dt
#from dateutil import parser as dtp
#import scipy.linalg as sla
#import scipy.signal as ssig
#import scipy.ndimage as ndi
import scipy.optimize as opti
#import scipy.interpolate as stp
#import scipy.spatial.distance as ssd
import matplotlib.pyplot as plt
#import matplotlib.cm as cm
#import matplotlib.ticker as mt
#import matplotlib._pylab_helpers as hlp
#from matplotlib.colors import LogNorm
#from matplotlib import colors
#import matplotlib.colors as mplcolors
#import matplotlib.gridspec as gridspec
from functools import partial
#from collections import OrderedDict
#import multiprocessing as mp
#np.set_printoptions(suppress=True, linewidth=160)
#import pandas as pd
#import statsmodels.api as sm
#import statsmodels.formula.api as smf
#from statsmodels.regression.quantile_regression import QuantReg
#import PIL.Image as pli
#import seaborn as sns
#import cmocean
#import theil_sen as ts
#import window_filter as wf
#import itertools as itt

##--------------------------------------------------------------------------##

## Home-brew robust statistics:
#try:
#    import robust_stats
#    reload(robust_stats)
#    rs = robust_stats
#except ImportError:
#    sys.stderr.write("\nError!  robust_stats module not found!\n"
#           "Please install and try again ...\n\n")
#    sys.exit(1)

## Fast FITS I/O:
#try:
#    import fitsio
#except ImportError:
#    sys.stderr.write("\nError: fitsio module not found!\n")
#    sys.exit(1)

## FITS I/O:
#try:
#    import astropy.io.fits as pf
#except ImportError:
#    try:
#       import pyfits as pf
#    except ImportError:
#        sys.stderr.write("\nError!  No FITS I/O module found!\n"
#               "Install either astropy.io.fits or pyfits and try again!\n\n")
#        sys.exit(1)

##--------------------------------------------------------------------------##
## Colors for fancy terminal output:
NRED    = '\033[0;31m'   ;  BRED    = '\033[1;31m'
NGREEN  = '\033[0;32m'   ;  BGREEN  = '\033[1;32m'
NYELLOW = '\033[0;33m'   ;  BYELLOW = '\033[1;33m'
NBLUE   = '\033[0;34m'   ;  BBLUE   = '\033[1;34m'
NMAG    = '\033[0;35m'   ;  BMAG    = '\033[1;35m'
NCYAN   = '\033[0;36m'   ;  BCYAN   = '\033[1;36m'
NWHITE  = '\033[0;37m'   ;  BWHITE  = '\033[1;37m'
ENDC    = '\033[0m'

## Suppress colors in cron jobs:
if (os.getenv('FUNCDEF') == '--nocolors'):
    NRED    = ''   ;  BRED    = ''
    NGREEN  = ''   ;  BGREEN  = ''
    NYELLOW = ''   ;  BYELLOW = ''
    NBLUE   = ''   ;  BBLUE   = ''
    NMAG    = ''   ;  BMAG    = ''
    NCYAN   = ''   ;  BCYAN   = ''
    NWHITE  = ''   ;  BWHITE  = ''
    ENDC    = ''

## Fancy text:
degree_sign = u'\N{DEGREE SIGN}'

## Dividers:
halfdiv = '-' * 40
fulldiv = '-' * 80

##--------------------------------------------------------------------------##
## Save FITS image with clobber (astropy / pyfits):
#def qsave(iname, idata, header=None, padkeys=1000, **kwargs):
#    this_func = sys._getframe().f_code.co_name
#    sys.stderr.write("Writing to '%s' ... " % iname)
#    if header:
#        while (len(header) < padkeys):
#            header.append() # pad header
#    if os.path.isfile(iname):
#        os.remove(iname)
#    pf.writeto(iname, idata, header=header, **kwargs)
#    sys.stderr.write("done.\n")

##--------------------------------------------------------------------------##
## Save FITS image with clobber (fitsio):
#def qsave(iname, idata, header=None, **kwargs):
#    this_func = sys._getframe().f_code.co_name
#    sys.stderr.write("Writing to '%s' ... " % iname)
#    #if os.path.isfile(iname):
#    #    os.remove(iname)
#    fitsio.write(iname, idata, clobber=True, header=header, **kwargs)
#    sys.stderr.write("done.\n")

##--------------------------------------------------------------------------##
##--------------------------------------------------------------------------##

## Mayhem extraction tools:
import nres_extraction
reload(nres_extraction)
nrex = nres_extraction
trio = nres_extraction.TraceIO()
frox = nres_extraction.FlatRelativeOptimalExtraction()
#test_traces = "./scratch/trace1_med_nres01_tung01_20181125_07d_06b_00a.fits"
test_traces = "derp.fits"
double_path = "images/med_nres01_thar01_20181125_07d_06b_00a.fits.fz"
for item in (test_traces, double_path):
    if not os.path.isfile(item):
        sys.stderr.write("Can't find file: '%s'\n" % item)
        sys.exit(1)

## Load traces for fiddling:
trdata = trio.load_traces(test_traces)

## Spectrograph/optics brilliance:
import spectrograph_optics
reload(spectrograph_optics)
import glass
reload(glass)

##--------------------------------------------------------------------------##
## Spectrograph hardware config items (constants):
nres_pixel_size_mm = 0.015  # 15-micron pixels
nres_gratio = 4.0           # NRES uses R4 grating
nres_ruling_lmm = 41.59     # lines per mm ruled
nres_blaze_angle_rad = np.arctan(nres_gratio)
nres_spacing_um = 1e3 / nres_ruling_lmm     # grating spacing in microns

nres_prism_glass = "PBM2"   # glass type used in cross-dispersing prism
nres_prism_apex_deg = 55.0  # apex angle of cross-dispersing prism
nres_prism_apex_rad = np.radians(nres_prism_apex_deg)

nres_center_wl_um = 0.479   # [I THINK] light wavelength nearest CCD center

nres_nominal_gamma = 0.0    # [radians] gamma angle for nres_center_wl_um
#nres_nominal_gamma = np.radians(1.0)

nres_focallen_mm = 375.15   # approximate camera focal length

##--------------------------------------------------------------------------##
sog = glass.Glass(nres_prism_glass)
nrp = spectrograph_optics.Prism(nres_prism_glass, nres_prism_apex_deg)
ogt = spectrograph_optics.GratingTools(nres_gratio,
        lines_per_mm=nres_ruling_lmm)

## Refractive index of central wavelength:
center_wl_nn = sog.refraction_index(nres_center_wl_um)

## Nominal incidence angle onto prism (~same for all wavelengths):
incident_ang_rad = np.arcsin(center_wl_nn * np.sin(0.5 * nres_prism_apex_rad))

## Minimum deflection angle at nominal wavelength:
min_dev_rad = 2.0 * incident_ang_rad - nres_prism_apex_rad

## Effective gamma angle is the nominal gamma angle plus the changes in prism
## deflection angle due to difference from nominal wavelength.
def gamma_eff(gamma_nom, wlen_um):
    # delta_gamma = deflection_rad - min_dev_rad
    #deflect_r = spectrograph_optics.prism_deflection_n(incident_ang_rad,
    #        nres_prism_apex_rad, sog.refraction_index(wlen_um))
    deflect_r = nrp.deflection_rad_wl(incident_ang_rad, wlen_um)
    return gamma_nom + deflect_r - min_dev_rad

use_gamma_eff = partial(gamma_eff, nres_nominal_gamma)

## From existing solutions ...
nres_sine_alpha = 0.971747764900

## To get central wavelengths, want to solve:
## m * λ = 2 * d * cos(γ(λ)) * cos(θ) * sin(θ_B)
## where
##  * θ   = facet illumination angle
##  * θ_B = blaze angle
#blaze_angle_rad = np.arctan(4.0)    # R4 grating
#gamma_angle_rad = 0.0               # grating gamma angle
#facet_angle_rad = np.radians(2.0)   # how close to blaze angle
#alpha_angle_rad = blaze_angle_rad + facet_angle_rad
blaze_r = np.arctan(4.0)
#facet_r = np.radians(0.5)
facet_r = np.arcsin(nres_sine_alpha) - blaze_r
fixed_geometry = 2.0 * nres_spacing_um * np.cos(facet_r) * np.sin(blaze_r)

def lamcen_residual(wlen, order=0):
    ls = wlen * order
    rs = fixed_geometry * np.cos(use_gamma_eff(wlen))
    return ls - rs

def iter_calc_lamcen(order):
    kw = {'order':order}
    runme = partial(lamcen_residual, **kw)
    return opti.bisect(runme, 0.0, 10.0)

## NOTES:
## alpha = blaze_angle + facet_angle
##  beta = blaze_angle - facet_angle    (center)

##--------------------------------------------------------------------------##
## Orders to evaluate:
useful_orders = 51.0 + np.arange(69.0)

## Old way:
spec_order_list = np.copy(useful_orders)
spec_order_wlmid, spec_order_FSR, spec_order_angsize = \
        ogt.get_order_params(spec_order_list, units='um')
spec_order_table = {kk:vv for kk,vv in zip(spec_order_list, spec_order_wlmid)}

## New way:
iter_order_wlmid = np.array([iter_calc_lamcen(x) for x in spec_order_list])
iter_order_table = {kk:vv for kk,vv in zip(spec_order_list, iter_order_wlmid)}

## Comparison of methods:
for ii,ww in enumerate(spec_order_wlmid):
    w2 = iter_order_wlmid[ii]
    sys.stderr.write("oid %3d --> %10.5f nm || %10.5f\n"
            % (ii, 1e3 * ww, 1e3 * w2))

##--------------------------------------------------------------------------##
## Given a set of central wavelengths, compute order separation:
def calc_order_sep(central_wlen_um):
    refr_idx = sog.refraction_index(central_wlen_um)
    #deflections = spectrograph_optics.prism_deflection_n(incident_ang_rad,
    #            nres_prism_apex_rad, refr_idx)
    deflections = nrp.deflection_rad_nr(incident_ang_rad, refr_idx)
    angle_change = deflections - min_dev_rad
    return 2*angle_change

nres_flen_pix = nres_focallen_mm / nres_pixel_size_mm
old_sep = calc_order_sep(spec_order_wlmid) * nres_flen_pix
new_sep = calc_order_sep(iter_order_wlmid) * nres_flen_pix

##--------------------------------------------------------------------------##
##--------------------------------------------------------------------------##
##--------------------------------------------------------------------------##
##--------------------------------------------------------------------------##
## Test of even better code:
dppgp = spectrograph_optics.DoublePassPrismGratingPrism()
lam_cen_dppgp = dppgp.calc_central_wlen_um(spec_order_list)

ctr_wlen, ctr_gamma = dppgp.fancy_deflections(spec_order_list)

#defl1_0, rbeam_0, defl2_0, h2_0 = dppgp.two_pass_deflection(ctr_wlen[0])
#defl1_z, rbeam_z, defl2_z, h2_z = dppgp.two_pass_deflection(ctr_wlen[-1])
ctr_headings, pg_yshifts, pc_yshifts = \
        np.array([dppgp.two_pass_deflection(x) for x in ctr_wlen]).T
dp_yshifts = pg_yshifts + pc_yshifts

##--------------------------------------------------------------------------##
##--------------------------------------------------------------------------##
##--------------------------------------------------------------------------##
##--------------------------------------------------------------------------##

## Store results for external use/comparison:
sys.stderr.write("Storing central wavelengths ... ")
with open('lam_ctr_data.csv', 'w') as f:
    f.write("oidx,spec_ord,lam_dumb,lam_iter,sep_dumb,sep_iter\n")
    norm_old_sep = old_sep - old_sep.min()
    norm_new_sep = new_sep - new_sep.min()
    for ii,w1 in enumerate(spec_order_wlmid):
        so = spec_order_list[ii]
        w2 = iter_order_wlmid[ii]
        s1, s2 = norm_old_sep[ii], norm_new_sep[ii]
        f.write("%2d,%3d,%8.6f,%8.6f,%8.3f,%8.3f\n"
                % (ii, so, w1, w2, s1, s2))
sys.stderr.write("done.\n")

##--------------------------------------------------------------------------##
##--------------------------------------------------------------------------##
##--------------------------------------------------------------------------##
##--------------------------------------------------------------------------##
##--------------------------------------------------------------------------##
## Dispersion is a function of angles and wavelength:
## dβ/dλ = n / (d * cosγ * cosβ)
## dβ/dλ = (sinα + sinβ) / (λ * cosβ)
#dbeta_dlambda = (nres_sine_alpha +

##--------------------------------------------------------------------------##
## Arrays of spectrograph coordinates:
ny, nx = 4096, 4096
x_list = (0.5 + np.arange(nx)) / nx - 0.5            # relative (centered)
y_list = (0.5 + np.arange(ny)) / ny - 0.5            # relative (centered)
#xx, yy = np.meshgrid(x_list, y_list)                 # relative (centered)
xx, yy = np.meshgrid(nx*x_list, ny*y_list)           # absolute (centered)
#xx, yy = np.meshgrid(np.arange(nx), np.arange(ny))   # absolute
#yy, xx = np.meshgrid(np.arange(ny), np.arange(nx), indexing='ij') # absolute
#yy, xx = np.nonzero(np.ones_like(img_vals))          # absolute
#yy, xx = np.mgrid[0:ny,   0:nx].astype('uint16')     # absolute (array)
#yy, xx = np.mgrid[1:ny+1, 1:nx+1].astype('uint16')   # absolute (pixel)
#xx *= nres_pixel_size_mm
#yy *= nres_pixel_size_mm

gamma2d = np.arcsin(yy * nres_pixel_size_mm / nres_focallen_mm)
##gamma2d = np.arctan(yy * nres_pixel_size_mm / nres_focallen_mm)
beta_2d = np.arcsin(xx * nres_pixel_size_mm / nres_focallen_mm) # beta - beta_c
del xx, yy, x_list, y_list

#ordernum = 100
#wlen2d_test = nres_spacing_um *

##--------------------------------------------------------------------------##
## Plot config:

# gridspec examples:
# https://matplotlib.org/users/gridspec.html

#gs1 = gridspec.GridSpec(4, 4)
#gs1.update(wspace=0.025, hspace=0.05)  # set axis spacing

#ax1 = plt.subplot2grid((3, 3), (0, 0), colspan=3) # top-left + center + right
#ax2 = plt.subplot2grid((3, 3), (1, 0), colspan=2) # mid-left + mid-center
#ax3 = plt.subplot2grid((3, 3), (1, 2), rowspan=2) # mid-right + bot-right
#ax4 = plt.subplot2grid((3, 3), (2, 0))            # bot-left
#ax5 = plt.subplot2grid((3, 3), (2, 1))            # bot-center

##--------------------------------------------------------------------------##
fig_dims = (12, 10)
fig = plt.figure(1, figsize=fig_dims)
plt.gcf().clf()
#fig, axs = plt.subplots(2, 2, sharex=True, figsize=fig_dims, num=1)
# sharex='col' | sharex='row'
#fig.frameon = False # disable figure frame drawing
#fig.subplots_adjust(left=0.07, right=0.95)
#ax1 = plt.subplot(gs[0, 0])
#ax1 = fig.add_subplot(111)
#ax1 = fig.add_axes([0, 0, 1, 1])
#ax1.patch.set_facecolor((0.8, 0.8, 0.8))
#ax1.grid(True)
#ax1.axis('off')

## Disable axis offsets:
#ax1.xaxis.get_major_formatter().set_useOffset(False)
#ax1.yaxis.get_major_formatter().set_useOffset(False)

#ax1.plot(kde_pnts, kde_vals)

#blurb = "some text"
#ax1.text(0.5, 0.5, blurb, transform=ax1.transAxes)
#ax1.text(0.5, 0.5, blurb, transform=ax1.transAxes,
#      va='top', ha='left', bbox=dict(facecolor='white', pad=10.0))
#      fontdict={'family':'monospace'}) # fixed-width

#colors = cm.rainbow(np.linspace(0, 1, len(plot_list)))
#for camid, c in zip(plot_list, colors):
#    cam_data = subsets[camid]
#    xvalue = cam_data['CCDATEMP']
#    yvalue = cam_data['PIX_MED']
#    yvalue = cam_data['IMEAN']
#    ax1.scatter(xvalue, yvalue, color=c, lw=0, label=camid)

#mtickpos = [2,5,7]
#ndecades = 1.0   # for symlog, set width of linear portion in units of dex
#nonposx='mask' | nonposx='clip' | nonposy='mask' | nonposy='clip'
#ax1.set_xscale('log', basex=10, nonposx='mask', subsx=mtickpos)
#ax1.set_xscale('log', nonposx='clip', subsx=[3])
#ax1.set_yscale('symlog', basey=10, linthreshy=0.1, linscaley=ndecades)
#ax1.xaxis.set_major_formatter(formatter) # re-format x ticks
#ax1.set_ylim(ax1.get_ylim()[::-1])
#ax1.set_xlabel('whatever', labelpad=30)  # push X label down

#ax1.set_xticks([1.0, 3.0, 10.0, 30.0, 100.0])
#ax1.set_xticks([1, 2, 3], ['Jan', 'Feb', 'Mar'])
#for label in ax1.get_xticklabels():
#    label.set_rotation(30)
#    label.set_fontsize(14)

#ax1.set_xlim(nice_limits(xvec, pctiles=[1,99], pad=1.2))
#ax1.set_ylim(nice_limits(yvec, pctiles=[1,99], pad=1.2))

#spts = ax1.scatter(x, y, lw=0, s=5)
#cbar = fig.colorbar(spts, orientation='vertical')
#cbar.formatter.set_useOffset(False)
#cbar.update_ticks()

#fig.tight_layout() # adjust boundaries sensibly, matplotlib v1.1+
#plt.draw()
#fig.savefig(plot_name, bbox_inches='tight')

# cyclical colormap ... cmocean.cm.phase
# cmocean: https://matplotlib.org/cmocean/

######################################################################
# CHANGELOG (compare_order_Y_positions.py):
#---------------------------------------------------------------------
#
#  2019-01-20:
#     -- Increased __version__ to 0.0.1.
#     -- First created compare_order_Y_positions.py.
#

ext: py

max_stars_repo_name: rsiverd/mayhem

max_issues_repo_name: rsiverd/mayhem

max_forks_repo_name: rsiverd/mayhem

reason: import numpy,from numpy,import scipy,import statsmodels,from statsmodels,import astropy


text:
import os
import tempfile
from io import StringIO
from json import dump
from zipfile import ZipFile

import numpy as np

from matplotlib import pyplot as plt

from astropy.time import Time
from astropy.table import Table
from astropy import units as u

from astropy.visualization import quantity_support

from aas_timeseries.backports import time_support
from aas_timeseries.colors import auto_assign_colors
from aas_timeseries.views import BaseView, View
from aas_timeseries.matplotlib import (PhaseAsDegreesLocator,
                                       PhaseAsDegreesFormatter,
                                       PhaseAsRadiansLocator,
                                       PhaseAsRadiansFormatter)

__all__ = ['InteractiveTimeSeriesFigure']

class InteractiveTimeSeriesFigure(BaseView):
    """
    An interactive time series figure.

    Parameters
    ----------
    width : int, optional
        The preferred width of the figure, in pixels.
    height : int, optional
        The preferred height of the figure, in pixels.
    padding : float, optional
        The padding inside the axes, as a fraction of the size of the axes
    resize : bool, optional
        Whether to resize the figure to the available space.
    title : str, optinal
        If views are added to the figure, this title is used for the default
        view, otherwise 'Default' is used.
    """

    def __init__(self, width=600, height=400, padding=36, resize=False, title=None, time_mode=None):
        super().__init__(time_mode=time_mode)
        self._width = width
        self._height = height
        self._resize = resize
        self._padding = padding
        self._yunit = 'auto'
        self._views = []
        self._title = title

    @property
    def yunit(self):
        """
        The unit to use for the y axis, or 'auto' to guess. If using data
        without quantities, this can be set to ``u.one``.
        """
        return self._yunit

    @yunit.setter
    def yunit(self, value):
        if isinstance(value, u.UnitBase) or value == 'auto':
            self._yunit = value
        else:
            self._yunit = u.Unit(value)

    def _guess_yunit(self):
        """
        Try and guess unit to use on the y axis based on the limits set and the
        data used.
        """

        # First check if the limits were set explicitly and if so use those
        # units if they were Quantities since it seems likely this was
        # deliberate.
        if self.ylim is not None and isinstance(self.ylim[0], u.Quantity):
            return self.ylim[0].unit

        # Next, check through all the data required by the layers, first in the
        # figure, then in the views, and return the first unit or lack of unit
        # found.
        for layer in self._layers:
            for (data, colname) in layer._required_ydata:
                return data.unit(colname)
        for view in self._views:
            for layer in view['view']._layers:
                for (data, colname) in layer._required_ydata:
                    return data.unit(colname)

        # Since we didn't find anything, let's assume that the y axis should
        # be dimensionless. In theory we could also go through non-data layers
        # and check the units used there.
        return u.one

    def add_view(self, title, description=None, include=None, exclude=None, empty=False, time_mode=None):

        if empty:
            inherited_layers = {}
        elif include is not None:
            for layer in include:
                if layer not in self._layers:
                    raise ValueError(f'Layer {layer} does not exist in base figure')
            inherited_layers = {layer: self._layers[layer] for layer in include}
        elif exclude is not None:
            for layer in exclude:
                if layer not in self._layers:
                    raise ValueError(f'Layer {layer} does not exist in base figure')
            inherited_layers = {layer: self._layers[layer] for layer in self._layers if layer not in exclude}
        else:
            inherited_layers = self._layers.copy()

        view = View(figure=self, inherited_layers=inherited_layers, time_mode=time_mode)

        self._views.append({'title': title, 'description': description, 'view': view})

        return view

    def export_interactive_bundle(self, filename, embed_data=False,
                                  minimize_data=True, override_style=False):
        """
        Create a bundle for the interactive figure containing an HTML file and
        JSON file along with any required CSV files.

        Parameters
        ----------
        filename : str
            The filename for the zip file
        embed_data : bool, optional
            Whether to embed the data in the JSON file (`True`) or include it
            in separate CSV files (`False`). The default is `False`.
        minimize_data : bool, optional
            Whether to include only data required for the visualization (`True`)
            or also other unused columns/fields in the time series (`False`).
            The default is `True`.
        override_style : bool, optional
            By default, any unspecified colors will be automatically chosen.
            If this parameter is set to `True`, all colors will be reassigned,
            even if already set.
        """

        start_dir = os.path.abspath('.')
        tmp_dir = tempfile.mkdtemp()
        os.chdir(tmp_dir)
        try:
            self.save_vega_json('figure.json', embed_data=embed_data)
        finally:
            os.chdir(start_dir)
        html_file = os.path.join(os.path.dirname(__file__), 'screenshot', 'template.html')
        with ZipFile(filename, 'w') as fzip:
            for filename in os.listdir(tmp_dir):
                fzip.write(os.path.join(tmp_dir, filename), os.path.basename(filename))
            fzip.write(html_file, 'index.html')

    def _check_colors(self, override_style=False):
        # Auto-assign colors if needed
        colors = auto_assign_colors(self._layers)
        for layer, color in zip(self._layers, colors):
            if override_style or layer.color is None:
                layer.color = color

    def save_static(self, prefix, format='png', override_style=False):
        """
        Export the figure to one or more static files using Matplotlib. If views
        are present then one plot is produced for each view.

        Parameters
        ----------
        prefix : str
            The name of the plot (without extension).
        format : str
            Any valid format supported by Matplotlib.
        override_style : bool, optional
            By default, any unspecified colors will be automatically chosen.
            If this parameter is set to `True`, all colors will be reassigned,
            even if already set.
        """

        # Start off by figuring out what units we are using on the y axis.
        # Note that we check the consistency of the units only here for
        # simplicity otherwise any guessing while users add/remove layers is
        # tricky.
        yunit = self._guess_yunit() if self.yunit == 'auto' else self.yunit

        # Auto-assign colors if needed
        self._check_colors(override_style=override_style)

        # We now loop over the main figure and all the views, and produce a
        # static plot for each of them.

        def pad_limits(limits, padding):
            vrange = (limits[1] - limits[0]) * padding
            return limits[0] - vrange, limits[1] + vrange

        for iview, view in enumerate([self] + self._views):

            if view is not self:
                view = view['view']

            # Note that if we aren't dealing with non-absolute times, the
            # following settings don't matter since the time_support context
            # manager won't have any effect.

            if view.time_format == 'auto' or view.time_mode != 'absolute':
                time_format = 'iso'
                simplify = True
            else:
                time_format = view.time_format
                simplify = False

            with time_support(format=time_format, simplify=simplify, scale='utc'):
                with quantity_support():

                    fig = plt.figure(figsize=(self._width / 100,
                                              self._height / 100))
                    ax = fig.add_axes([0.15, 0.12, 0.8, 0.86])

                    for layer in view.layers:
                        layer.to_mpl(ax, yunit=yunit)

            if view.time_mode == 'phase':
                if view.time_format == 'degrees':
                    ax.xaxis.set_major_locator(PhaseAsDegreesLocator())
                    ax.xaxis.set_major_formatter(PhaseAsDegreesFormatter())
                elif view.time_format == 'radians':
                    ax.xaxis.set_major_locator(PhaseAsRadiansLocator())
                    ax.xaxis.set_major_formatter(PhaseAsRadiansFormatter())

            x_domain, y_domain = view._get_domains(yunit, as_vega=False)

            ax.set_xlim(*x_domain)
            ax.set_ylim(*y_domain)

            # Apply padding - we just get the limits again because the x limits
            # above may have been Time objects, so we get the limits again from
            # Matplotlib.
            ax.set_xlim(*pad_limits(ax.get_xlim(), self._padding / self._width))
            ax.set_ylim(*pad_limits(ax.get_ylim(), self._padding / self._height))

            if view is self:
                filename = prefix + '.' + format
            else:
                filename = prefix + '_view' + str(iview) + '.' + format

            ax.set_xlabel(view.xlabel)
            ax.set_ylabel(view.ylabel)

            fig.savefig(filename)

    def save_vega_json(self, filename, embed_data=False,
                       minimize_data=True, override_style=False):
        """
        Export the JSON file, and optionally CSV data files.

        Parameters
        ----------
        filename : str
            The name of the JSON file
        embed_data : bool, optional
            Whether to embed the data in the JSON file (`True`) or include it
            in separate CSV files (`False`). The default is `False`.
        minimize_data : bool, optional
            Whether to include only data required for the visualization (`True`)
            or also other unused columns/fields in the time series (`False`).
            The default is `True`.
        override_style : bool, optional
            By default, any unspecified colors will be automatically chosen.
            If this parameter is set to `True`, all colors will be reassigned,
            even if already set.
        """

        # Start off by figuring out what units we are using on the y axis.
        # Note that we check the consistency of the units only here for
        # simplicity otherwise any guessing while users add/remove layers is
        # tricky.
        yunit = self._guess_yunit() if self.yunit == 'auto' else self.yunit

        # Auto-assign colors if needed
        self._check_colors(override_style=override_style)

        # Start off with empty JSON
        json = {}

        # Schema
        json['$schema'] = 'https://vega.github.io/schema/vega/v4.json'

        # Layout
        json['title'] = self._title or 'Default'
        json['width'] = self._width
        json['height'] = self._height
        json['padding'] = 0
        json['autosize'] = {'type': 'fit', 'resize': self._resize}

        json['_extend'] = {}

        # Data

        # We start off by checking which columns and data are going to be
        # required. We do this by iterating over the layers in the main
        # figure and the views and keeping track of the set of (data, column)
        # that are needed.
        required_xdata = []
        required_ydata = []
        required_tooltipdata = []
        for layer in self._layers:
            required_xdata.extend(layer._required_xdata)
            required_ydata.extend(layer._required_ydata)
            required_tooltipdata.extend(layer._required_tooltipdata)
        for view in self._views:
            for layer in view['view']._layers:
                required_xdata.extend(layer._required_xdata)
                required_ydata.extend(layer._required_ydata)
                required_tooltipdata.extend(layer._required_tooltipdata)
        required_xdata = set(required_xdata)
        required_ydata = set(required_ydata)
        required_tooltipdata = set(required_tooltipdata)

        json['data'] = []

        for data in self._data.values():

            # Start off by constructing a new table with only the subset of
            # columns required, and the time as an ISO string. Note that we
            # need to explicitly specify that we want UTC times, then add the
            # Z suffix since this isn't something that astropy does. For
            # relative times we always use seconds, and for phases we use values
            # in the range [0:1].
            table = Table()
            time_columns = []
            for colname in data.time_series.colnames:
                if (not minimize_data or (data, colname) in required_xdata | required_ydata | required_tooltipdata):
                    column = data.time_series[colname]
                    if isinstance(column, Time):
                        table[colname] = np.char.add(column.utc.isot, 'Z')
                        time_columns.append(colname)
                    elif (data, colname) in required_xdata:
                        try:
                            table[colname] = data.column_to_values(colname, u.s)
                        except u.UnitsError:
                            table[colname] = data.column_to_values(colname, u.one)
                    elif (data, colname) in required_ydata:
                        table[colname] = data.column_to_values(colname, yunit)
                    else:
                        table[colname] = column

            # Next up, we create the information for the 'parse' Vega key
            # which indicates the format for each column.
            parse = {}
            for colname in table.colnames:
                column = table[colname]
                if colname in time_columns:
                    parse[colname] = 'date'
                elif column.dtype.kind in 'fi':
                    parse[colname] = 'number'
                elif column.dtype.kind in 'b':
                    parse[colname] = 'boolean'
                else:
                    parse[colname] = 'string'

            vega = {'name': data.uuid,
                    'format': {'type': 'csv',
                               'parse': parse}}

            # We now either embed the CSV inside the JSON or create CSV files.
            # For now we use UUIDs for the data file names in the latter case
            # but in future we could find a way to preserve information about
            # the original filenames the data came from.
            if embed_data:
                s = StringIO()
                table.write(s, format='ascii.basic', delimiter=',')
                s.seek(0)
                # NOTE: when embedding the data inside the JSON file, we should
                # just use simple Unix line endings inside the serialized table.
                csv_string = s.read().replace('\r\n', '\n')
                vega['values'] = csv_string
            else:
                data_filename = 'data_' + data.uuid + '.csv'
                data_path = os.path.join(os.path.dirname(filename), data_filename)
                table.write(data_path, format='ascii.basic', delimiter=',')
                vega['url'] = data_filename

            json['data'].append(vega)

        # At this point, we loop over all the views (including the main view
        # given by self) and output these to the JSON.

        for view in [self] + self._views:

            if view is self:

                view_json = json

            else:

                view_json = {'name': self.uuid,
                             'title': view['title'],
                             'description': view['description']}

                if '_views' not in json:
                    json['_views'] = []

                json['_views'].append(view_json)

                view = view['view']

            if view._time_mode == 'absolute':
                x_type = 'time'
                x_input = 'iso'
            elif view._time_mode == 'relative':
                x_type = 'number'
                x_input = 'seconds'
            elif view._time_mode == 'phase':
                x_type = 'number'
                x_input = 'unity'

            view_json['_extend'] = {'scales': [{'name': 'xscale',
                                                'input': x_input,
                                                'output': view.time_format}]}

            view_json['axes'] = [{'orient': 'bottom',
                                  'scale': 'xscale',
                                  'title': view.xlabel},
                                 {'orient': 'left',
                                  'scale': 'yscale',
                                  'title': view.ylabel}]

            view_json['scales'] = [{'name': 'xscale',
                                    'type': x_type,
                                    'range': 'width',
                                    'zero': False,
                                    'padding': self._padding},
                                   {'name': 'yscale',
                                    'type': 'log' if view.ylog else 'linear',
                                    'range': 'height',
                                    'zero': False,
                                    'padding': self._padding}]

            # Limits, if specified

            x_domain, y_domain = view._get_domains(yunit)

            if x_domain is not None:
                view_json['scales'][0]['domain'] = x_domain

            if y_domain is not None:
                view_json['scales'][1]['domain'] = y_domain

            if view is self:

                view_json['marks'] = []
                for layer, settings in self._layers.items():
                    view_json['marks'].extend(layer.to_vega(yunit=yunit))

            else:

                view_json['markers'] = []

                for layer, settings in view._inherited_layers.items():
                    for uuid in layer.uuids:
                        view_json['markers'].append({'name': uuid,
                                                     'visible': settings['visible']})

                for layer, settings in view._layers.items():

                    if 'marks' not in json['_extend']:
                        json['_extend']['marks'] = []

                    json['_extend']['marks'].extend(layer.to_vega(yunit=yunit))
                    for uuid in layer.uuids:
                        view_json['markers'].append({'name': uuid,
                                                     'visible': settings['visible']})

        with open(filename, 'w') as f:
            dump(json, f, indent='  ', sort_keys=True)

    def preview_interactive(self):
        """
        Show an interactive version of the figure (only works in Jupyter
        notebook or lab).
        """
        # FIXME: should be able to do without a file
        from jupyter_aas_timeseries import TimeSeriesWidget
        tmpfile = tempfile.mktemp()
        self.save_vega_json(tmpfile, embed_data=True, minimize_data=True)
        widget = TimeSeriesWidget(tmpfile)
        return widget

    def remove(self, layer):
        """
        Remove a layer from the figure.
        """
        if layer in self._layers:
            self._layers.pop(layer)
            for view in self._views:
                if layer in view['view'].layers:
                    view['view'].remove(layer)
        else:
            raise ValueError(f"Layer '{layer.label}' is not in figure")

max_stars_repo_name: astrofrog/aas-time-series-affiliated

max_issues_repo_name: astrofrog/aas-time-series-affiliated

max_forks_repo_name: astrofrog/aas-time-series-affiliated

reason: import numpy,from astropy


text:
# -*- coding: utf-8 -*-
"""
Implementation of the Dueling Posterior Sampling algorithm with Bayesian
logistic regression credit assignment.
"""

import numpy as np
from collections import defaultdict
from scipy.optimize import minimize

from Learning_algorithms.DPS_helper_functions import advance, get_state_action_visit_counts

def DPS_log_reg(time_horizon, hyper_params, env, num_iter, diri_prior = 1,
            run_str = '', seed = None):
    """
    This function implements the DPS algorithm with a Bayesian logistic
    regression credit assignment model over state/action rewards.

    Inputs:
        1) time_horizon: episode horizon; this is the number of state/action
           pairs in each learning episode.
        2) hyper_params: the hyperparameters for the logistic regression credit
           assignment model. This is a length-2 list of the form
           [prior_covariance, cov_scale], where the prior covariance and
           covariance scale are both scalars.
        3) env: the RL environment.
        4) num_iter: the number of iterations of the learning algorithm to run.
           Note that two trajectory rollouts occur per iteration of learning.
        5) diri_prior: parameter for setting the prior of the transition
           dynamics model. For each state/action pair, the Dirichlet prior is
           set to diri_prior * np.ones(num_states), where num_states is the
           number of states in the MDP.
        6) run_str: if desired, a string with information about the current
           call to DPS_log_reg (e.g. hyperparameter values or repetition number),
           which can be useful for print statements to track progress.
        7) seed: seed for random number generation.

    Returns: a vector of rewards received as the algorithm runs. This is either
             a) the total rewards from each trajectory rollout, or b) the
             rewards at every step taken in the environment (the environment
             determines whether a) or b) is used).
    """

    # Unpack hyperparameters:
    [prior_covariance, cov_scale] = hyper_params

    if not seed is None:
        np.random.seed(seed)

    # Numbers of states and actions in the environment:
    num_states = env.nS
    num_actions = env.nA

    num_sa_pairs = num_states * num_actions   # Number of state/action pairs

    # Initialize prior mean and covariance for Bayesian logistic regression model:
    prior_mean = np.zeros(num_sa_pairs)
    prior_cov = prior_covariance * np.eye(num_sa_pairs)

    # Prior model, without covariance scaling:
    LR_prior_model = {'mean': prior_mean, 'cov': prior_cov}

    # Eigenvalues and eigenvectors of the prior covariance matrix:
    evals, evecs = np.linalg.eigh(prior_cov)

    # Initially, reward model is just the prior, since we don't have any data.
    LR_model = {'mean': prior_mean, 'cov_evecs': evecs, \
                'cov_evals': cov_scale * evals}

    # Dirichlet model posterior over state/action transition probabilities.
    # Initially, this is set to the Dirichlet prior, and it's updated after
    # each observed state transition. Note that dirichlet_posterior[state][action]
    # is a length-num_states array, specifying the probability distribution for
    # transitioning to each possible subsequent state from the given state/action.
    # Setting diri_prior = 1 gives a uniform prior over transition probabilities.
    dirichlet_posterior = defaultdict(lambda: defaultdict(lambda: \
                                            diri_prior * np.ones(num_states)))

    # For each trajectory pair, store difference between how many times each
    # state/action pair is visited:
    observation_matrix = np.empty((0, num_sa_pairs))

    # Preference labels corresponding to the observations:
    preference_labels = np.empty((0, 1))

    num_policies = 2     # Number of policies to sample per learning iteration

    # To store results (for evaluation purposes only):
    if env.store_episode_reward:   # Store total reward for each trajectory
        rewards = np.empty(num_iter * num_policies)
    else:    # Store reward at each step within each trajectory
        rewards = np.empty(num_iter * time_horizon * num_policies)

    reward_count = 0   # Counts how many values in the "rewards" variable
                       # defined above have been populated

    """
    Here is where the learning algorithm begins.
    """
    for iteration in range(num_iter):

        # Print status:
        print('Bayesian logistic regression, parameters %s: iteration = %i' % \
              (run_str, iteration + 1))

        # Sample policies:
        policies = advance(num_policies, dirichlet_posterior, LR_model,
                        num_states, num_actions, time_horizon)

        # Roll out trajectories using these policies:
        trajectories = []

        for policy in policies:    # Roll out 2 action sequences

            state = env.reset()

            state_sequence = np.empty(time_horizon + 1)
            action_sequence = np.empty(time_horizon)

            for t in range(time_horizon):

                action = np.random.choice(num_actions, p = policy[t, state, :])

                next_state, done = env.step(action)

                state_sequence[t] = state
                action_sequence[t] = action

                # Tracking rewards for evaluation purposes (in case of
                # tracking rewards at every single step):
                if not env.store_episode_reward:
                    rewards[reward_count] = env.get_step_reward(state,
                                action, next_state)
                    reward_count += 1

                # Terminate trajectory if environment turns on "done" flag.
                if done:
                    state_sequence = state_sequence[: t + 2]
                    action_sequence = action_sequence[: t + 1]

                    break

                # Update state transition posterior:
                dirichlet_posterior[state][action][next_state] += 1

                state = next_state

            state_sequence[-1] = next_state
            trajectories.append([state_sequence, action_sequence])

            # Tracking rewards for evaluation purposes (in case of tracking
            # rewards just over entire episodes):
            if env.store_episode_reward:

                rewards[reward_count] = env.get_trajectory_return()
                reward_count += 1

        # Obtain a preference between the 2 trajectories:
        preference = env.get_trajectory_preference(trajectories[0],
                    trajectories[1])

        # This only matters if using deterministic preferences (in this case,
        # there can be a tie between two trajectories. In this case, we skip
        # updating the reward posterior):
        if preference == 0.5:
            continue

        # Store the difference in state/action visitation counts for the 2
        # trajectories:
        visitations = []

        for trajectory in trajectories:

            visitation_vec = get_state_action_visit_counts(trajectory,
                                        num_states, num_actions)

            visitations.append(visitation_vec)

        visitation_diff = visitations[1] - visitations[0]

        observation_matrix = np.vstack((observation_matrix, visitation_diff))

        # Store preference information. Preference gets mapped to 0.5 if
        # 2nd trajectory is preferred, and -0.5 otherwise.
        preference_labels = np.vstack((preference_labels,
                        np.reshape([preference - 0.5], (1, 1))))

        # Call feedback function to update the reward model posterior by
        # performing credit assignment via Bayesian logistic regression:
        LR_model = feedback_logistic(LR_prior_model, observation_matrix,
                            preference_labels, cov_scale)

    # Return performance results:
    return rewards

def feedback_logistic(LR_prior_model, observation_matrix, preference_labels,
                      cov_scale, r_init = []):
    """
    This function updates the posterior over rewards based on the new
    preference data, via Bayesian logistic regression credit assignment.
    The Bayesian logistic regression posterior is calculated via the Laplace
    approximation.

    Inputs (note: d is the number of state/action pairs; n is the number of
            data points/observations/rows in the observation matrix):
        1) LR_prior_model: the prior (without covariance scaling via cov_scale),
           represented as a dictionary with keys 'mean' and 'cov'. These are
           the prior mean (length-d array) and prior covariance (d-by-d matrix)
           respectively.
        2) observation_matrix: n-by-d array, in which each row corresponds
           to an observation.
        3) preference_labels: length-n vector, in which each element is the
           label corresponding to an observation.
        4) cov_scale: covariance scaling hyperparameter; a positive scalar.
        5) (Optional) r_init: initial guess for convex optimization; length-d
           NumPy array when specified.

    Output:
        The updated model posterior, represented as a dictionary with keys
        'mean', 'cov_evecs', and 'cov_evals'. 'mean' is the posterior mean, a
        length-d NumPy array in which each element corresponds to a state/action
        pair. 'cov_evecs' is an d-by-d NumPy array in which each column is an
        eigenvector of the posterior covariance, and 'cov_evals' is a length-d
        array of the eigenvalues of the posterior covariance.
    """

    # Unpack the prior
    prior_mean = LR_prior_model['mean']
    prior_cov_inv = np.linalg.inv(LR_prior_model['cov'])  # Invert covariance

    # Number of state/action pairs:
    num_sa_pairs = observation_matrix.shape[1]

    # Solve convex optimization problem to obtain the posterior mean reward
    # vector via the Laplace approximation:
    if r_init == []:
        r_init = 0.5 * np.ones(num_sa_pairs)    # Initial guess

    res = minimize(logreg_objective, r_init, args = (observation_matrix,
                   preference_labels, prior_mean, prior_cov_inv),
                   method = 'L-BFGS-B', jac = logreg_obj_gradient)

    post_mean = res.x

    # Approximate inverse of the posterior covariance by evaluating the
    # objective function's Hessian at the MAP estimate:
    post_cov_inverse = logreg_obj_hessian(post_mean, observation_matrix,
                                  preference_labels, prior_cov_inv)

    # Calculate the eigenvectors and eigenvalues of the inverse posterior
    # covariance matrix:
    evals, evecs = np.linalg.eigh(post_cov_inverse)

    # Invert the eigenvalues to get the eigenvalues corresponding to the
    # covariance matrix:
    evals = cov_scale / evals

    # Return the model posterior:
    return {'mean': post_mean, 'cov_evecs': evecs, 'cov_evals': evals}

def logreg_objective(r, X, y, r0, prior_cov_inv):
    """
    Evaluate the optimization objective function for finding the posterior
    mean of the Bayesian logistic regression model via the Laplace
    approximation; the posterior mean is the minimum of this (convex) objective
    function.

    Inputs (note: d is the number of state/action pairs; n is the number of
            data points/observations/rows in the observation matrix):
        1) r: the "point" at which to evaluate the objective function. This is
           a length-d vector.
        2) X: n-times-d observation matrix.
        3) y: length-n vector of labels (all values should be +/- 1)
        4) r0: prior value for r (length-d vector)
        5) prior_cov_inv: d x d inverse of the prior covariance matrix

    Output: the objective function evaluated at the given point (r).
    """

    obj = 0.5 * (r - r0) @ prior_cov_inv @ (r - r0)

    for i in range(len(y)):

        obj += np.log(1 + np.exp(-y[i] * np.dot(X[i, :], r)))

    return obj

def logreg_obj_gradient(r, X, y, r0, prior_cov_inv):
    """
    Evaluate the gradient of the Laplace approximation objective function.

    Inputs: same as in logreg_objective.

    Output: the objective function's gradient evaluated at the given point (r).
    """

    grad = prior_cov_inv @ (r - r0)

    for i in range(len(y)):

        yi = y[i]; xi = X[i, :]

        grad -= yi * xi / (1 + np.exp(yi * np.dot(xi, r)))

    return grad

def logreg_obj_hessian(r, X, y, prior_cov_inv):
    """
    Evaluate the Hessian of the Laplace approximation objective function.

    Inputs: same as in logreg_GP_objective, except without r0.

    Output: the objective function's Hessian matrix evaluated at the given
            point (r).
    """

    hess = prior_cov_inv

    for i in range(len(y)):

        yi = y[i]; xi = X[i, :]
        xi_vec = xi.reshape((len(xi), 1))

        hess += (xi_vec @ np.transpose(xi_vec)) * np.exp(yi * np.dot(xi, r)) /   \
                (np.exp(yi * np.dot(xi, r)) + 1)**2

    return hess

max_stars_repo_name: ernovoseller/DuelingPosteriorSampling

max_issues_repo_name: ernovoseller/DuelingPosteriorSampling

max_forks_repo_name: ernovoseller/DuelingPosteriorSampling

reason: import numpy,from scipy


text:
import argparse
import glob
from os import PathLike
from pathlib import Path
from subprocess import Popen, PIPE
import sys
from typing import List, Union

import numpy as np
import pandas as pd

from schrodinger import structure
from schrodinger.structutils import analyze, measure

def lig_avg_bfact(lig:analyze.Ligand) -> float:
    '''
    Calculates the average b-factor (temperature factor) across a ligand.

    Args:
    lig - ligand object from schrodinger
    Returns:
    float value of average b-factor across the ligand
    '''
    #deal with protons on some ligands that have Nonetype temperature factors with None check
    return np.average(np.array([a.temperature_factor for a in lig.st.atom if a.temperature_factor is not None]))

def get_chains_around_lig(cplx: structure.Structure, lig: analyze.Ligand, distance: float = 5.0) -> List[structure._Chain]:
    '''
    Calls measure to get chain object within distance A of ligand

    Args:
    cplx - the structure object holding everything in the PDB file
    lig - the ligand object of interest
    distance - maximum distance of atoms to consider

    Returns:
    List of schrodinger.structure._Chain objects that have atoms within distance Angstrom of the ligand
    '''
    atms = measure.get_atoms_close_to_structure(cplx, lig.st, distance)
    chs = list(set([cplx.atom[i].chain for i in atms]))
    return chs

def download_pdbs(base_dir: Union[str, PathLike], pdb_id_list: List[str]) -> None:

    '''
    Downloads a set of PDB IDs into base_dir/raw/ID/ID.pdb

    Args:
    base_dir - data directory for this project
    pdb_id_list - list of RCSB PDB IDs to download

    Returns:
    None
    '''

    base = Path(base_dir)
    raws = base / 'raw'

    if not raws.exists():
        raws.mkdir(parents=True)

    for pdb in pdb_id_list:
        #create a directory for each file
        outpath = raws / (str(pdb))
        outpath.mkdir(parents=True, exist_ok=False)

        pdb_url = f"https://files.rcsb.org/download/{pdb}.pdb"
        command = f"wget {pdb_url} -P {outpath}"

        pc = Popen(command, stdout= PIPE, stderr= PIPE, shell=True)
        o, e = pc.communicate()

        if pc.returncode != 0:
            print(f"Error with {pdb} : {e.decode('utf-8')}")

def convert_pdb_to_maestro(pdb_file: Union[str, PathLike]) -> None:
    '''
    Converts a PDB file to a Maestro File using Schrodinger's structconvert script

    Args:
    pdb_file - Filename or pathlike object of the pdb file to convert

    Returns:
    None
    '''

    pdbf = Path(pdb_file)
    maestro_file = pdbf.with_suffix(".mae")
    #will throw a deprecated warning about -ipdb and -omae since newest versions infer filetype through extension. Keeping as-is
    #since it still works and should be compatible with previous versions
    command = f"$SCHRODINGER/utilities/structconvert -ipdb {pdbf} -omae {maestro_file}"

    pc = Popen(command, stdout= PIPE, stderr= PIPE, shell=True)
    o, e = pc.communicate()

    if pc.returncode != 0:
        print(f"Error with {pdb_file}: {e.decode('utf-8')}")

def process_structure(filename: Union[str, PathLike], ligname: str) -> None:
    '''
    Pre-processes the raw Maestro File for downstream applications
    This script will identify and keep chains that are in close proximity to the ligand of interest, and will change the ligand's
    chain to 'L' for ease of downstream processing.
    The output is a processed file of the form filename_proc.mae

    Args:
    filename - maestro file to be processed
    ligname - ligand's 3-letter code

    Returns:
    None
    '''
    alpha = list('ABCDEFGHIJKMNOPQRSTUVWXYZ')

    fin = Path(filename)

    #read in maestro file
    with structure.StructureReader(filename) as st:
        cplx = next(st)

    real_ligs = [l for l in analyze.find_ligands(cplx) if l.pdbres.strip() == ligname]

    #future fix - some structures fail the find_ligands command, need to figure out why
    if not real_ligs:
        print(f"\nError with {fin.stem}")
    else:
        #sort by b factor if there are multiple copies of the ligand, take the lig with lowest b-factor
        real_ligs.sort(key=lambda x: lig_avg_bfact(x))
        lig_to_take = real_ligs[0]

        chs = get_chains_around_lig(cplx, lig_to_take)

        #remove the ligand from the structure, still have a copy in ligs_to_take
        ##TODO: Remove the indices of all ligands that are not the primary ligand. Just extend the atom indexes for all ligands in real_ligs

        l_atoms = lig_to_take.atom_indexes
        cplx.deleteAtoms(l_atoms)

        #make a new structure of the chains around the ligand, rename alphabetically, but don't use chain L. Ligand will be chain L
        chains = [cplx.chain[c] for c in chs]
        for i,c in enumerate(chains):
            c.name = alpha[i]

        #create a new structure object with these chains
        newstruc = chains[0].extractStructure()
        if len(chains) > 1:
            for c in chains[1:]:
                newstruc.extend(c.extractStructure())

        #rename ligand
        ligst = lig_to_take.st
        lig_chain_name = list(ligst.chain)[0].name
        ligst.chain[lig_chain_name].name = 'L'
        #add ligand to new structure
        newstruc.extend(ligst)

        #write to outfile
        outfile = fin.parent / (fin.stem + '_proc.mae')
        newstruc.write(str(outfile))

if __name__ == "__main__":

    parser = argparse.ArgumentParser(description="Process Maestro Files")
    parser.add_argument("-o", "--output_dir", help="Directory to save files to")
    parser.add_argument("-s", help="csv of structures to download")

    args = parser.parse_args()

    pdb_info = pd.read_csv(str(args.s), header=0, index_col=0)
    pdbs = list(pdb_info.rscb_id)
    #download pdb files
    download_pdbs(args.output_dir, pdbs)
    #convert them to Maestro
    for pdb_file in glob.glob(f"{args.output_dir}/raw/*/*.pdb"):
        convert_pdb_to_maestro(pdb_file)
    #process Maestro to keep chains around ligand, and rename ligand to chain 'L'
    for i,f in enumerate(glob.glob(f"{args.output_dir}/raw/*/*.mae")):
        fpath = Path(f)
        pdb_name = fpath.stem

        if '_proc' in pdb_name:
            #already processed file
            continue
        #because tqdm is for some reason disabled in schrodinger's python... use an old school progress reporter.
        sys.stdout.write(f"\rWorking on {i}/{len(pdb_info)} | {pdb_name}")
        sys.stdout.flush()

        info_df = pdb_info[pdb_info['rscb_id'] == str(pdb_name)]

        l3letter = np.array(info_df['lig_3letter'])[0]

        process_structure(f, l3letter)

max_stars_repo_name: songlab-cal/contact-geometry

max_issues_repo_name: songlab-cal/contact-geometry

max_forks_repo_name: songlab-cal/contact-geometry

reason: import numpy