Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Issue/cloest pv systenms #103

Closed
wants to merge 3 commits into from
Closed
Changes from 1 commit
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
Next Next commit
add option to pick closets pv systems, not random + tests
  • Loading branch information
peterdudfield committed Nov 17, 2022
commit 8ea85d0c6b514082084d87d71487807116360e1e
Original file line number Diff line number Diff line change
@@ -1,6 +1,8 @@
"""Ensure there is N PV systems per example"""
import logging

from typing import Optional

import numpy as np
import xarray as xr
from torchdata.datapipes import functional_datapipe
@@ -13,29 +15,61 @@
class EnsureNPVSystemsPerExampleIterDataPipe(IterDataPipe):
"""Ensure there is N PV systems per example"""

def __init__(self, source_datapipe: IterDataPipe, n_pv_systems_per_example: int, seed=None):
def __init__(
self,
source_datapipe: IterDataPipe,
n_pv_systems_per_example: int,
seed=None,
method: str = "random",
locations_datapipe: Optional[IterDataPipe] = None,
):
"""
Ensure there is N PV systems per example

Args:
source_datapipe: Datapipe of PV data
n_pv_systems_per_example: Number of PV systems to have in example
seed: Random seed for choosing
method: method for picking PV systems. Can be 'random' or 'closest'
locations_datapipe: location of this example. Can be None as its only needed for 'closest'
"""
self.source_datapipe = source_datapipe
self.n_pv_systems_per_example = n_pv_systems_per_example
self.rng = np.random.default_rng(seed=seed)
self.method = method
self.locations_datapipe = locations_datapipe

assert method in ["random", "closest"]

if method == "closest":
assert (
locations_datapipe is not None
), f"If you are slect closest PV systems, then a location data pipe is needed"

def __iter__(self):
for xr_data in self.source_datapipe:
if len(xr_data.pv_system_id) > self.n_pv_systems_per_example:
logger.debug(f"Reducing PV systems to {self.n_pv_systems_per_example}")
# More PV systems are available than we need. Reduce by randomly sampling:
subset_of_pv_system_ids = self.rng.choice(
xr_data.pv_system_id,
size=self.n_pv_systems_per_example,
replace=False,
)
if self.method == "random":
subset_of_pv_system_ids = self.rng.choice(
xr_data.pv_system_id,
size=self.n_pv_systems_per_example,
replace=False,
)
elif self.method == "closest":

location = next(self.locations_datapipe)

# get distance
delta_x = xr_data.x_osgb - location.x
delta_y = xr_data.y_osgb - location.y
r2 = delta_x ** 2 + delta_y ** 2

# order and select closest
r2 = r2.sortby(r2)
subset_of_pv_system_ids = r2.pv_system_id[: self.n_pv_systems_per_example]

xr_data = xr_data.sel(pv_system_id=subset_of_pv_system_ids)
elif len(xr_data.pv_system_id) < self.n_pv_systems_per_example:
logger.debug("Padding out PV systems")
57 changes: 57 additions & 0 deletions tests/transform/xarray/pv/test_ensure_n_pv_systems_per_example.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,57 @@
import numpy as np
import pytest

from ocf_datapipes.transform.xarray import EnsureNPVSystemsPerExample
from ocf_datapipes.utils.consts import Location


def test_ensure_n_pv_systems_per_example_expand(passiv_datapipe):

data_before = next(iter(passiv_datapipe))

passiv_datapipe = EnsureNPVSystemsPerExample(passiv_datapipe, n_pv_systems_per_example=12)
data_after = next(iter(passiv_datapipe))

assert len(data_before[0, :]) == 2
assert len(data_after[0, :]) == 12


def test_ensure_n_pv_systems_per_example_random(passiv_datapipe):

data_before = next(iter(passiv_datapipe))

passiv_datapipe = EnsureNPVSystemsPerExample(passiv_datapipe, n_pv_systems_per_example=1)
data_after = next(iter(passiv_datapipe))

assert len(data_before[0, :]) == 2
assert len(data_after[0, :]) == 1


def test_ensure_n_pv_systems_per_example_closest_error(passiv_datapipe):

with pytest.raises(Exception):
_ = EnsureNPVSystemsPerExample(
passiv_datapipe, n_pv_systems_per_example=1, method="closest"
)


def test_ensure_n_pv_systems_per_example_closest(passiv_datapipe):

# make fake location datapipe
location = Location(x=2.687e05, y=6.267e05)
location_datapipe = iter([location])

data_before = next(iter(passiv_datapipe))

passiv_datapipe = EnsureNPVSystemsPerExample(
passiv_datapipe,
n_pv_systems_per_example=1,
method="closest",
locations_datapipe=location_datapipe,
)

data_after = next(iter(passiv_datapipe))

assert len(data_before[0, :]) == 2
assert len(data_after[0, :]) == 1
assert data_after.pv_system_id[0] == 9960