polytope-python 2.0.1__tar.gz → 2.0.2__tar.gz

{polytope_python-2.0.1/polytope_python.egg-info → polytope_python-2.0.2}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: polytope-python
-Version: 2.0.1
+Version: 2.0.2
 Summary: Polytope datacube feature extraction library
 Author-email: "European Centre for Medium-Range Weather Forecasts (ECMWF)" <software.support@ecmwf.int>
 Maintainer-email: James Hawkes <James.Hawkes@ecmwf.int>, Mathilde Leuridan <Mathilde.Leuridan@ecmwf.int>

{polytope_python-2.0.1 → polytope_python-2.0.2}/polytope_feature/_version.py

@@ -1,2 +1,2 @@
 # Do not change! Do not track in version control!
-__version__ = "2.0.1"
+__version__ = "2.0.2"

{polytope_python-2.0.1 → polytope_python-2.0.2}/polytope_feature/datacube/transformations/datacube_type_change/datacube_type_change.py

@@ -78,6 +78,30 @@ class TypeChangeStrToInt(DatacubeAxisTypeChange):
         return tuple(values)
 
 
+class TypeChangeStrToFloat(DatacubeAxisTypeChange):
+    def __init__(self, axis_name, new_type):
+        self.axis_name = axis_name
+        self._new_type = new_type
+
+    def transform_type(self, value):
+        try:
+            return float(value)
+        except ValueError:
+            return None
+
+    def make_str(self, value):
+        values = []
+        for val in value:
+            int_val = int(val)
+            if val == int_val:
+                # we return a str of the integer value
+                values.append(str(int_val))
+            else:
+                # we return a str of the float value
+                values.append(str(val))
+        return tuple(values)
+
+
 class TypeChangeStrToTimestamp(DatacubeAxisTypeChange):
     def __init__(self, axis_name, new_type):
         self.axis_name = axis_name
@@ -195,6 +219,7 @@ _type_to_datacube_type_change_lookup = {
     "int": "TypeChangeStrToInt",
     "date": "TypeChangeStrToTimestamp",
     "time": "TypeChangeStrToTimedelta",
+    "float": "TypeChangeStrToFloat",
     "subhourly_step_compact": "TypeChangeSubHourlyTimeStepsCompact",
     "subhourly_step": "TypeChangeSubHourlyTimeSteps",
 }

{polytope_python-2.0.1 → polytope_python-2.0.2}/polytope_feature/engine/quadtree_slicer.py

@@ -26,9 +26,16 @@ class QuadTreeSlicer(Engine):
 
     def extract_single(self, datacube, polytope):
         # extract a single polygon
+        # if need to find nearest points, then take alternative slicing method using quadtree to find nearest point
         if use_rust:
-            polytope_points = [tuple(point) for point in polytope.points]
-            polygon_points = self.quad_tree.query_polygon(self.points, 0, polytope_points)
+            if len(datacube.nearest_search) == 0:
+                polytope_points = [tuple(point) for point in polytope.points]
+                polygon_points = self.quad_tree.query_polygon(self.points, 0, polytope_points)
+            else:
+                nn_points = [tuple(pt) for pt in datacube.nearest_search[tuple(polytope.axes())]]
+                polygon_points = []
+                for nn_pt in nn_points:
+                    polygon_points.append(self.quad_tree.nearest_neighbor(nn_pt, self.points))
         else:
             polygon_points = self.quad_tree.query_polygon(polytope)
         return polygon_points

{polytope_python-2.0.1 → polytope_python-2.0.2}/polytope_feature/polytope.py

@@ -10,7 +10,7 @@ from .engine.optimised_quadtree_slicer import OptimisedQuadTreeSlicer
 from .engine.point_in_polygon_slicer import PointInPolygonSlicer
 from .engine.quadtree_slicer import QuadTreeSlicer
 from .options import PolytopeOptions
-from .shapes import ConvexPolytope, Product
+from .shapes import ConvexPolytope, Point, Product, Union
 from .utility.combinatorics import group, tensor_product
 from .utility.exceptions import AxisOverdefinedError
 from .utility.list_tools import unique
@@ -122,6 +122,7 @@ class Polytope:
 
     def slice(self, datacube, polytopes: List[ConvexPolytope]):
         """Low-level API which takes a polytope geometry object and uses it to slice the datacube"""
+
         self.find_compressed_axes(datacube, polytopes)
 
         self.remove_compressed_axis_in_union(polytopes)
@@ -176,17 +177,37 @@ class Polytope:
         slicer_type = self.engine_options[ax.name]
         return self.engines[slicer_type]
 
+    def switch_polytope_dim(self, request):
+        # If we see a 2-dim slicer on an axis
+        # then make sure that if the shape is a point, we set decompose_1D to False
+        for ax, slicer in self.engine_options.items():
+            if slicer == "quadtree":
+                for shp in request.shapes:
+                    if ax in shp.axes() and isinstance(shp, Point):
+                        shp.decompose_1D = False
+                    elif isinstance(shp, Union):
+                        for s in shp._shapes:
+                            if ax in s.axes() and isinstance(s, Point):
+                                s.decompose_1D = False
+
     def retrieve(self, request: Request, method="standard"):
         """Higher-level API which takes a request and uses it to slice the datacube"""
         logging.info("Starting request for %s ", self.context)
         self.datacube.check_branching_axes(request)
+        self.switch_polytope_dim(request)
         for polytope in request.polytopes():
             method = polytope.method
             if method == "nearest":
-                if self.datacube.nearest_search.get(tuple(polytope.axes()), None) is None:
-                    self.datacube.nearest_search[tuple(polytope.axes())] = polytope.values
+                if polytope.is_flat:
+                    if self.datacube.nearest_search.get(tuple(polytope.axes()), None) is None:
+                        self.datacube.nearest_search[tuple(polytope.axes())] = polytope.values
+                    else:
+                        self.datacube.nearest_search[tuple(polytope.axes())].append(polytope.values[0])
                 else:
-                    self.datacube.nearest_search[tuple(polytope.axes())].append(polytope.values[0])
+                    if self.datacube.nearest_search.get(tuple(polytope.axes()), None) is None:
+                        self.datacube.nearest_search[tuple(polytope.axes())] = polytope.points
+                    else:
+                        self.datacube.nearest_search[tuple(polytope.axes())].append(polytope.points[0])
         request_tree = self.slice(self.datacube, request.polytopes())
         logging.info("Created request tree for %s ", self.context)
         self.datacube.get(request_tree, self.context)

{polytope_python-2.0.1 → polytope_python-2.0.2}/polytope_feature/shapes.py

@@ -77,8 +77,12 @@ class Product(Shape):
         assert len(self._axes) == len(all_axes)
 
         self._polytopes = []
+        self.is_flat = True
         for poly in polytopes:
             self._polytopes.extend(poly.polytope())
+            for p in poly.polytope():
+                if len(p.axes()) > 1:
+                    self.is_flat = False
 
         self.is_in_union = False
         self.method = method
@@ -128,6 +132,7 @@ class Point(Shape):
         self._axes = axes
         self.values = values
         self.method = method
+        self.decompose_1D = True
         assert len(values) == 1
 
     def axes(self):
@@ -137,11 +142,15 @@ class Point(Shape):
         # TODO: change this to use the Product instead and return a Product here of the two 1D selects
 
         polytopes = []
-        for point in self.values:
-            poly_to_mult = []
-            for i in range(len(self._axes)):
-                poly_to_mult.append(ConvexPolytope([self._axes[i]], [[point[i]]], self.method, is_orthogonal=True))
-            polytopes.append(Product(*poly_to_mult, method=self.method, value=[point]))
+        if self.decompose_1D:
+            for point in self.values:
+                poly_to_mult = []
+                for i in range(len(self._axes)):
+                    poly_to_mult.append(ConvexPolytope([self._axes[i]], [[point[i]]], self.method, is_orthogonal=True))
+                polytopes.append(Product(*poly_to_mult, method=self.method, value=[point]))
+        else:
+            for point in self.values:
+                polytopes.append(ConvexPolytope(self._axes, [point], self.method, is_orthogonal=True))
         self.polytopes = polytopes
 
         return self.polytopes
@@ -464,19 +473,18 @@ class Union(Shape):
         self._axes = axes
         for s in shapes:
             assert s.axes() == self.axes()
-
-        self.polytopes = []
         self._shapes = shapes
 
-        for s in shapes:
-            for poly in s.polytope():
-                poly.add_to_union()
-                self.polytopes.append(poly)
-
     def axes(self):
         return self._axes
 
     def polytope(self):
+        self.polytopes = []
+
+        for s in self._shapes:
+            for poly in s.polytope():
+                poly.add_to_union()
+                self.polytopes.append(poly)
         return self.polytopes
 
     def __repr__(self):

{polytope_python-2.0.1 → polytope_python-2.0.2/polytope_python.egg-info}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: polytope-python
-Version: 2.0.1
+Version: 2.0.2
 Summary: Polytope datacube feature extraction library
 Author-email: "European Centre for Medium-Range Weather Forecasts (ECMWF)" <software.support@ecmwf.int>
 Maintainer-email: James Hawkes <James.Hawkes@ecmwf.int>, Mathilde Leuridan <Mathilde.Leuridan@ecmwf.int>

{polytope_python-2.0.1 → polytope_python-2.0.2}/polytope_python.egg-info/SOURCES.txt

@@ -177,6 +177,7 @@ polytope_python.egg-info/dependency_links.txt
 polytope_python.egg-info/requires.txt
 polytope_python.egg-info/top_level.txt
 rust/Cargo.toml
+rust/src/distance.rs
 rust/src/healpix_nested.rs
 rust/src/lambert_conformal.rs
 rust/src/lib.rs
@@ -246,6 +247,7 @@ tests/test_request_tree.py
 tests/test_request_trees_after_slicing.py
 tests/test_reverse_transformation.py
 tests/test_shapes.py
+tests/test_shapes_volume.py
 tests/test_slice_date_range_fdb.py
 tests/test_slice_date_range_fdb_v2.py
 tests/test_slice_fdb_box.py

polytope_python-2.0.2/rust/src/distance.rs

@@ -0,0 +1,33 @@
+pub fn dist2(a: (f64, f64), b: (f64, f64)) -> f64 {
+    let dx = a.0 - b.0;
+    let dy = a.1 - b.1;
+    dx * dx + dy * dy
+}
+
+pub fn box_dist2(center: (f64, f64), size: (f64, f64), point: (f64, f64)) -> f64 {
+    let half_w = size.0 / 2.0;
+    let half_h = size.1 / 2.0;
+
+    let min_x = center.0 - half_w;
+    let max_x = center.0 + half_w;
+    let min_y = center.1 - half_h;
+    let max_y = center.1 + half_h;
+
+    let dx = if point.0 < min_x {
+        min_x - point.0
+    } else if point.0 > max_x {
+        point.0 - max_x
+    } else {
+        0.0
+    };
+
+    let dy = if point.1 < min_y {
+        min_y - point.1
+    } else if point.1 > max_y {
+        point.1 - max_y
+    } else {
+        0.0
+    };
+
+    dx * dx + dy * dy
+}

{polytope_python-2.0.1 → polytope_python-2.0.2}/rust/src/lib.rs

@@ -3,6 +3,7 @@ pub mod lambert_conformal;
 pub mod list_tools;
 use pyo3::wrap_pyfunction;
 
+pub mod distance;
 use crate::lambert_conformal::{get_latlons_oblate, get_latlons_sphere};
 
 pub mod healpix_nested;

{polytope_python-2.0.1 → polytope_python-2.0.2}/rust/src/quadtree_mod.rs

@@ -9,6 +9,7 @@ use pyo3::exceptions::PyRuntimeError;
 // TODO: look at rust built in arena
 
 use crate::slicing_tools::{is_contained_in, slice_in_two};
+use crate::distance::{dist2, box_dist2};
 
 
 
@@ -54,6 +55,16 @@ impl QuadTree {
         }
     }
 
+    fn nearest_neighbor(&self, query: (f64, f64), quadtree_points: Vec<(f64, f64)>) -> Option<usize> {
+        if self.nodes.is_empty() {
+            return None;
+        }
+        let mut best_idx = None;
+        let mut best_dist2 = f64::INFINITY;
+        self.nn_search(0, query, &mut best_idx, &mut best_dist2, &quadtree_points);
+        best_idx
+    }
+
     fn sizeof(&self) -> usize {
         let mut size = size_of::<Self>();
         let nodes_size: usize = self.nodes.len() * size_of::<QuadTreeNode>();
@@ -145,6 +156,39 @@ impl QuadTree {
         points.into_iter().map(|(x, y)| [x,y]).collect()
     }
 
+    fn nn_search(
+        &self,
+        node_idx: usize,
+        query: (f64, f64),
+        best_idx: &mut Option<usize>,
+        best_dist2: &mut f64,
+        quadtree_points: &Vec<(f64, f64)>
+    ) {
+        let node = &self.nodes[node_idx];
+
+        // if this node is farther than the current best, ignore
+        if box_dist2(node.center, node.size, query) > *best_dist2 {
+            return;
+        }
+
+        // compare distance of points inside leaf node
+        if let Some(point_indices) = &node.points {
+            for &pi in point_indices {
+                let p = quadtree_points[pi];
+                let d2 = dist2(p, query);
+                if d2 < *best_dist2 {
+                    *best_dist2 = d2;
+                    *best_idx = Some(pi);
+                }
+            }
+            return;
+        }
+        // else, recurse into children
+        for &child_idx in &node.children {
+            self.nn_search(child_idx, query, best_idx, best_dist2, &quadtree_points);
+        }
+    }
+
     fn convert_polygon(&self, points: Option<Vec<(f64, f64)>>) -> Option<Vec<[f64; 2]>> {
         points.map(|pts| pts.into_iter().map(|(x, y)| [x, y]).collect())
     }

{polytope_python-2.0.1 → polytope_python-2.0.2}/tests/test_lambert_lam_grid_unstructured_fdb.py

@@ -6,7 +6,7 @@ import pytest
 from helper_functions import find_nearest_latlon
 
 from polytope_feature.polytope import Polytope, Request
-from polytope_feature.shapes import Box, Select
+from polytope_feature.shapes import Box, Point, Select, Union
 
 
 class TestQuadTreeSlicer:
@@ -111,3 +111,98 @@ class TestQuadTreeSlicer:
         # plt.scatter(eccodes_lons, eccodes_lats, s=6, c="green")
         # plt.colorbar(label="Temperature")
         # plt.show()
+
+    @pytest.mark.fdb
+    def test_quad_tree_slicer_extract_point(self):
+        import pygribjump as gj
+
+        request = Request(
+            Select("date", [pd.Timestamp("20250221T0000")]),
+            Select("step", [0]),
+            Select("param", ["130"]),
+            Select("levtype", ["sfc"]),
+            Point(["latitude", "longitude"], [[44.25, 5.55]], method="nearest"),
+        )
+
+        self.fdbdatacube = gj.GribJump()
+
+        self.API = Polytope(
+            datacube=self.fdbdatacube,
+            options=self.options,
+        )
+
+        result = self.API.retrieve(request)
+        assert len(result.leaves) == 1
+        result.pprint()
+
+        lats = []
+        lons = []
+        eccodes_lats = []
+        eccodes_lons = []
+        tol = 1e-3
+        leaves = result.leaves
+        for i in range(len(leaves)):
+            cubepath = leaves[i].flatten()
+            lat = cubepath["latitude"][0]
+            lon = cubepath["longitude"][0]
+            lats.append(lat)
+            lons.append(lon)
+            nearest_points = find_nearest_latlon("tests/data/lambert_lam_one_message.grib", lat, lon)
+            eccodes_lat = nearest_points[0][0]["lat"]
+            eccodes_lon = nearest_points[0][0]["lon"]
+            eccodes_lats.append(eccodes_lat)
+            eccodes_lons.append(eccodes_lon)
+            assert eccodes_lat - tol <= lat
+            assert lat <= eccodes_lat + tol
+            assert eccodes_lon - tol <= lon
+            assert lon <= eccodes_lon + tol
+
+    @pytest.mark.fdb
+    def test_quad_tree_slicer_extract_point_union(self):
+        import pygribjump as gj
+
+        pt1 = Point(["latitude", "longitude"], [[44.25, 5.55]], method="nearest")
+        pt2 = Point(["latitude", "longitude"], [[43.75, 5.35]], method="nearest")
+
+        pt_union = Union(["latitude", "longitude"], pt1, pt2)
+
+        request = Request(
+            Select("date", [pd.Timestamp("20250221T0000")]),
+            Select("step", [0]),
+            Select("param", ["130"]),
+            Select("levtype", ["sfc"]),
+            pt_union,
+        )
+
+        self.fdbdatacube = gj.GribJump()
+
+        self.API = Polytope(
+            datacube=self.fdbdatacube,
+            options=self.options,
+        )
+
+        result = self.API.retrieve(request)
+        assert len(result.leaves) == 2
+        result.pprint()
+
+        lats = []
+        lons = []
+        eccodes_lats = []
+        eccodes_lons = []
+        tol = 1e-3
+        leaves = result.leaves
+        for i in range(len(leaves)):
+            cubepath = leaves[i].flatten()
+            lat = cubepath["latitude"][0]
+            lon = cubepath["longitude"][0]
+            lats.append(lat)
+            lons.append(lon)
+            nearest_points = find_nearest_latlon("tests/data/lambert_lam_one_message.grib", lat, lon)
+            eccodes_lat = nearest_points[0][0]["lat"]
+            eccodes_lon = nearest_points[0][0]["lon"]
+            eccodes_lats.append(eccodes_lat)
+            eccodes_lons.append(eccodes_lon)
+            assert eccodes_lat - tol <= lat
+            assert lat <= eccodes_lat + tol
+            assert eccodes_lon - tol <= lon
+            assert lon <= eccodes_lon + tol

polytope_python-2.0.2/tests/test_shapes_volume.py

@@ -0,0 +1,21 @@
+from polytope_feature.shapes import ConvexPolytope, Point, Product
+
+
+class TestShapeDimension:
+    def setup_method(self, method):
+        pass
+
+    def test_point_dimension(self):
+        point_1D = Point(["lat", "lon"], [[1, 1]])
+
+        for shp in point_1D.polytope():
+            assert isinstance(shp, Product)
+            for polytope in shp._polytopes:
+                assert len(polytope.axes()) == 1
+
+        point_2D = Point(["lat", "lon"], [[1, 1]])
+        point_2D.decompose_1D = False
+
+        for shp in point_2D.polytope():
+            assert isinstance(shp, ConvexPolytope)
+            assert len(shp.axes()) == 2

{polytope_python-2.0.1 → polytope_python-2.0.2}/tests/test_type_change_transformation.py

@@ -3,6 +3,7 @@ import pandas as pd
 import xarray as xr
 
 from polytope_feature.datacube.transformations.datacube_type_change.datacube_type_change import (
+    TypeChangeStrToFloat,
     TypeChangeSubHourlyTimeSteps,
     TypeChangeSubHourlyTimeStepsCompact,
 )
@@ -33,6 +34,15 @@ class TestIntTypeChangeTransformation:
         result.pprint()
         assert result.leaves[0].flatten()["step"] == (0,)
 
+    def test_float_type_change_axis(self):
+        type_change_transform = TypeChangeStrToFloat("step", "float")
+
+        assert type_change_transform.transform_type("0.5") == 0.5
+        assert type_change_transform.transform_type("0") == 0.0
+
+        assert type_change_transform.make_str([0.1]) == ("0.1",)
+        assert type_change_transform.make_str([0.0]) == ("0",)
+
     def test_subhourly_step_type_change_axis(self):
         type_change_transform = TypeChangeSubHourlyTimeSteps("step", "subhourly_step")