diff --git a/notebooks/kdtree.ipynb b/notebooks/kdtree.ipynb
index 3e4b86a2e7abcbdfb00cdd54d4c5298a79336dba..8d1e6740d9a6b92995cb4c9b053f854404230753 100644
--- a/notebooks/kdtree.ipynb
+++ b/notebooks/kdtree.ipynb
@@ -1,5 +1,17 @@
{
"cells": [
+ {
+ "cell_type": "markdown",
+ "id": "f7143f08-1d06-4e94-bbf6-ef35ddd11556",
+ "metadata": {},
+ "source": [
+ "# KDTree\n",
+ "\n",
+ "Testing the time to look-up nearby records with the `KDTree` implementation. Note that this implementation is actually a `2DTree` since it can only compute a valid distance comparison between longitude and latitude positions.\n",
+ "\n",
+ "The `KDTree` object is used for finding the closest neighbour to a position, in this implementation we use the Haversine distance to compare positions."
+ ]
+ },
{
"cell_type": "code",
"execution_count": 1,
@@ -8,11 +20,9 @@
"outputs": [],
"source": [
"import os\n",
- "import gzip\n",
"os.environ[\"POLARS_MAX_THREADS\"] = \"4\"\n",
"\n",
- "from datetime import datetime, timedelta\n",
- "from random import choice\n",
+ "from datetime import datetime\n",
"from string import ascii_letters, digits\n",
"import random\n",
"import inspect\n",
@@ -20,7 +30,17 @@
"import polars as pl\n",
"import numpy as np\n",
"\n",
- "from GeoSpatialTools import Record, haversine, KDTree"
+ "from GeoSpatialTools import Record, KDTree"
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "id": "ec6c6e7f-8eee-47ea-a5e9-12537bb3412d",
+ "metadata": {},
+ "source": [
+ "## Set-up functions\n",
+ "\n",
+ "Used for generating data, or for comparisons by doing brute-force approach."
]
},
{
@@ -31,6 +51,7 @@
"outputs": [],
"source": [
"def randnum() -> float:\n",
+ " \"\"\"Get a random number between -1 and 1\"\"\"\n",
" return 2 * (np.random.rand() - 0.5)"
]
},
@@ -42,6 +63,7 @@
"outputs": [],
"source": [
"def generate_uid(n: int) -> str:\n",
+ " \"\"\"Generates a psuedo uid by randomly selecting from characters\"\"\"\n",
" chars = ascii_letters + digits\n",
" return \"\".join(random.choice(chars) for _ in range(n))"
]
@@ -49,6 +71,179 @@
{
"cell_type": "code",
"execution_count": 4,
+ "id": "9e647ecd-abdc-46a0-8261-aa081fda2e1d",
+ "metadata": {
+ "scrolled": true
+ },
+ "outputs": [],
+ "source": [
+ "def check_cols(\n",
+ " df: pl.DataFrame | pl.LazyFrame,\n",
+ " cols: list[str],\n",
+ " var_name: str = \"dataframe\",\n",
+ ") -> None:\n",
+ " \"\"\"\n",
+ " Check that a dataframe contains a list of columns. Raises an error if not.\n",
+ "\n",
+ " Parameters\n",
+ " ----------\n",
+ " df : polars Frame\n",
+ " Dataframe to check\n",
+ " cols : list[str]\n",
+ " Requried columns\n",
+ " var_name : str\n",
+ " Name of the Frame - used for displaying in any error.\n",
+ " \"\"\"\n",
+ " calling_func = inspect.stack()[1][3]\n",
+ " if isinstance(df, pl.DataFrame):\n",
+ " have_cols = df.columns\n",
+ " elif isinstance(df, pl.LazyFrame):\n",
+ " have_cols = df.collect_schema().names()\n",
+ " else:\n",
+ " raise TypeError(\"Input Frame is not a polars Frame\")\n",
+ "\n",
+ " cols_in_frame = intersect(cols, have_cols)\n",
+ " missing = [c for c in cols if c not in cols_in_frame]\n",
+ "\n",
+ " if len(missing) > 0:\n",
+ " err_str = f\"({calling_func}) - {var_name} missing required columns. \"\n",
+ " err_str += f'Require: {\", \".join(cols)}. '\n",
+ " err_str += f'Missing: {\", \".join(missing)}.'\n",
+ " raise ValueError(err_str)\n",
+ "\n",
+ " return\n",
+ "\n",
+ "\n",
+ "def haversine_df(\n",
+ " df: pl.DataFrame | pl.LazyFrame,\n",
+ " lon: float,\n",
+ " lat: float,\n",
+ " R: float = 6371,\n",
+ " lon_col: str = \"lon\",\n",
+ " lat_col: str = \"lat\",\n",
+ ") -> pl.DataFrame | pl.LazyFrame:\n",
+ " \"\"\"\n",
+ " Compute haversine distance on earth surface between lon-lat positions\n",
+ " in a polars DataFrame and a lon-lat position.\n",
+ "\n",
+ " Parameters\n",
+ " ----------\n",
+ " df : polars.DataFrame\n",
+ " The data, containing required columns:\n",
+ " * lon_col\n",
+ " * lat_col\n",
+ " * date_var\n",
+ " lon : float\n",
+ " The longitude of the position.\n",
+ " lat : float\n",
+ " The latitude of the position.\n",
+ " R : float\n",
+ " Radius of earth in km\n",
+ " lon_col : str\n",
+ " Name of the longitude column\n",
+ " lat_col : str\n",
+ " Name of the latitude column\n",
+ "\n",
+ " Returns\n",
+ " -------\n",
+ " polars.DataFrame\n",
+ " With additional column specifying distances between consecutive points\n",
+ " in the same units as 'R'. With colname defined by 'out_colname'.\n",
+ " \"\"\"\n",
+ " required_cols = [lon_col, lat_col]\n",
+ "\n",
+ " check_cols(df, required_cols, \"df\")\n",
+ " return (\n",
+ " df.with_columns(\n",
+ " [\n",
+ " pl.col(lat_col).radians().alias(\"_lat0\"),\n",
+ " pl.lit(lat).radians().alias(\"_lat1\"),\n",
+ " (pl.col(lon_col) - lon).radians().alias(\"_dlon\"),\n",
+ " (pl.col(lat_col) - lat).radians().alias(\"_dlat\"),\n",
+ " ]\n",
+ " )\n",
+ " .with_columns(\n",
+ " (\n",
+ " (pl.col(\"_dlat\") / 2).sin().pow(2)\n",
+ " + pl.col(\"_lat0\").cos()\n",
+ " * pl.col(\"_lat1\").cos()\n",
+ " * (pl.col(\"_dlon\") / 2).sin().pow(2)\n",
+ " ).alias(\"_a\")\n",
+ " )\n",
+ " .with_columns(\n",
+ " (2 * R * (pl.col(\"_a\").sqrt().arcsin()))\n",
+ " .round(2)\n",
+ " .alias(\"_dist\")\n",
+ " )\n",
+ " .drop([\"_lat0\", \"_lat1\", \"_dlon\", \"_dlat\", \"_a\"])\n",
+ " )\n",
+ "\n",
+ "\n",
+ "def intersect(a, b) -> set:\n",
+ " \"\"\"Intersection of a and b, items in both a and b\"\"\"\n",
+ " return set(a) & set(b)\n",
+ "\n",
+ "\n",
+ "def nearest_ship(\n",
+ " lon: float,\n",
+ " lat: float,\n",
+ " df: pl.DataFrame,\n",
+ " lon_col: str = \"lon\",\n",
+ " lat_col: str = \"lat\",\n",
+ ") -> pl.DataFrame:\n",
+ " \"\"\"\n",
+ " Find the observation nearest to a position in space.\n",
+ "\n",
+ " Get a frame with only the records that is closest to the input point.\n",
+ "\n",
+ " Parameters\n",
+ " ----------\n",
+ " lon : float\n",
+ " The longitude of the position.\n",
+ " lat : float\n",
+ " The latitude of the position.\n",
+ " df : polars.DataFrame\n",
+ " The pool of records to search. Can be pre-filtered and filter_datetime\n",
+ " set to False.\n",
+ " lon_col : str\n",
+ " Name of the longitude column in the pool DataFrame\n",
+ " lat_col : str\n",
+ " Name of the latitude column in the pool DataFrame\n",
+ "\n",
+ " Returns\n",
+ " -------\n",
+ " polars.DataFrame\n",
+ " Containing only records from the pool within max_dist of the input\n",
+ " point, optionally at the same datetime if filter_datetime is True.\n",
+ " \"\"\"\n",
+ " required_cols = [lon_col, lat_col]\n",
+ " check_cols(df, required_cols, \"df\")\n",
+ "\n",
+ " return (\n",
+ " df\n",
+ " .pipe(\n",
+ " haversine_df,\n",
+ " lon=lon,\n",
+ " lat=lat,\n",
+ " lon_col=lon_col,\n",
+ " lat_col=lat_col,\n",
+ " )\n",
+ " .filter(pl.col(\"_dist\").eq(pl.col(\"_dist\").min()))\n",
+ " .drop([\"_dist\"])\n",
+ " )\n"
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "id": "287bdc1d-1ecf-4c59-af95-d2dc639c6894",
+ "metadata": {},
+ "source": [
+ "## Initialise random data"
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": 5,
"id": "c60b30de-f864-477a-a09a-5f1caa4d9b9a",
"metadata": {},
"outputs": [
@@ -63,11 +258,11 @@
"│ --- ┆ --- │\n",
"│ i64 ┆ i64 │\n",
"â•žâ•â•â•â•â•â•â•ªâ•â•â•â•â•â•¡\n",
- "│ 127 ┆ 21 │\n",
- "│ -148 ┆ 36 │\n",
- "│ -46 ┆ -15 │\n",
- "│ 104 ┆ 89 │\n",
- "│ -57 ┆ -31 │\n",
+ "│ 62 ┆ -29 │\n",
+ "│ 146 ┆ 1 │\n",
+ "│ 104 ┆ 60 │\n",
+ "│ -162 ┆ -66 │\n",
+ "│ 72 ┆ 69 │\n",
"└──────┴─────┘\n"
]
}
@@ -76,7 +271,12 @@
"N = 16_000\n",
"lons = pl.int_range(-180, 180, eager=True)\n",
"lats = pl.int_range(-90, 90, eager=True)\n",
- "dates = pl.datetime_range(datetime(1900, 1, 1, 0), datetime(1900, 1, 31, 23), interval=\"1h\", eager=True)\n",
+ "dates = pl.datetime_range(\n",
+ " datetime(1900, 1, 1, 0),\n",
+ " datetime(1900, 1, 31, 23),\n",
+ " interval=\"1h\",\n",
+ " eager=True,\n",
+ ")\n",
"\n",
"lons_use = lons.sample(N, with_replacement=True).alias(\"lon\")\n",
"lats_use = lats.sample(N, with_replacement=True).alias(\"lat\")\n",
@@ -90,7 +290,7 @@
},
{
"cell_type": "code",
- "execution_count": 5,
+ "execution_count": 6,
"id": "875f2a67-49fe-476f-add1-b1d76c6cd8f9",
"metadata": {},
"outputs": [],
@@ -98,9 +298,19 @@
"records = [Record(**r) for r in df.rows(named=True)]"
]
},
+ {
+ "cell_type": "markdown",
+ "id": "bd83330b-ef2c-478e-9a7b-820454d198bb",
+ "metadata": {},
+ "source": [
+ "## Intialise the `KDTree`\n",
+ "\n",
+ "There is an overhead to constructing a `KDTree` object, so performance improvement is only for multiple comparisons."
+ ]
+ },
{
"cell_type": "code",
- "execution_count": 6,
+ "execution_count": 7,
"id": "1e883e5a-5086-4c29-aff2-d308874eae16",
"metadata": {},
"outputs": [
@@ -108,8 +318,8 @@
"name": "stdout",
"output_type": "stream",
"text": [
- "CPU times: user 151 ms, sys: 360 ms, total: 511 ms\n",
- "Wall time: 57.3 ms\n"
+ "CPU times: user 32.7 ms, sys: 1.4 ms, total: 34.1 ms\n",
+ "Wall time: 33.4 ms\n"
]
}
],
@@ -118,9 +328,30 @@
"kt = KDTree(records)"
]
},
+ {
+ "cell_type": "markdown",
+ "id": "0a37ef06-2691-4e01-96a9-1c1ecd582599",
+ "metadata": {},
+ "source": [
+ "## Compare with brute force approach"
+ ]
+ },
{
"cell_type": "code",
- "execution_count": 7,
+ "execution_count": 8,
+ "id": "365bbf30-7a93-438d-92b2-a3471f1e9249",
+ "metadata": {},
+ "outputs": [],
+ "source": [
+ "test_record = Record(\n",
+ " random.choice(range(-179, 180)) + randnum(),\n",
+ " random.choice(range(-89, 90)) + randnum(),\n",
+ ")"
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": 9,
"id": "69022ad1-5ec8-4a09-836c-273ef452451f",
"metadata": {},
"outputs": [
@@ -128,19 +359,18 @@
"name": "stdout",
"output_type": "stream",
"text": [
- "203 μs ± 4.56 μs per loop (mean ± std. dev. of 7 runs, 10,000 loops each)\n"
+ "130 μs ± 847 ns per loop (mean ± std. dev. of 7 runs, 10,000 loops each)\n"
]
}
],
"source": [
"%%timeit\n",
- "test_record = Record(random.choice(range(-179, 180)) + randnum(), random.choice(range(-89, 90)) + randnum())\n",
"kt.query(test_record)"
]
},
{
"cell_type": "code",
- "execution_count": 8,
+ "execution_count": 10,
"id": "28031966-c7d0-4201-a467-37590118e851",
"metadata": {},
"outputs": [
@@ -148,19 +378,45 @@
"name": "stdout",
"output_type": "stream",
"text": [
- "8.87 ms ± 188 μs per loop (mean ± std. dev. of 7 runs, 100 loops each)\n"
+ "8.34 ms ± 83.4 μs per loop (mean ± std. dev. of 7 runs, 100 loops each)\n"
]
}
],
"source": [
"%%timeit\n",
- "test_record = Record(random.choice(range(-179, 180)) + randnum(), random.choice(range(-89, 90)) + randnum())\n",
"np.argmin([test_record.distance(p) for p in records])"
]
},
{
"cell_type": "code",
- "execution_count": 9,
+ "execution_count": 11,
+ "id": "09e0f923-ca49-47bf-8643-e0b3a6d0467c",
+ "metadata": {},
+ "outputs": [
+ {
+ "name": "stdout",
+ "output_type": "stream",
+ "text": [
+ "8.28 ms ± 105 μs per loop (mean ± std. dev. of 7 runs, 100 loops each)\n"
+ ]
+ }
+ ],
+ "source": [
+ "%%timeit\n",
+ "nearest_ship(lon=test_record.lon, lat=test_record.lat, df=df)"
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "id": "f0359950-942d-45ea-8676-b22c8ce9e296",
+ "metadata": {},
+ "source": [
+ "## Verify that results are correct"
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": 12,
"id": "0d10b2ba-57b2-475c-9d01-135363423990",
"metadata": {},
"outputs": [
@@ -168,8 +424,8 @@
"name": "stdout",
"output_type": "stream",
"text": [
- "CPU times: user 17.4 s, sys: 147 ms, total: 17.6 s\n",
- "Wall time: 17.6 s\n"
+ "CPU times: user 16.9 s, sys: 144 ms, total: 17 s\n",
+ "Wall time: 17 s\n"
]
}
],
@@ -177,18 +433,28 @@
"%%time\n",
"n_samples = 1000\n",
"tol = 1e-8\n",
- "test_records = [Record(random.choice(range(-179, 180)) + randnum(), random.choice(range(-89, 90)) + randnum()) for _ in range(n_samples)]\n",
+ "test_records = [\n",
+ " Record(\n",
+ " random.choice(range(-179, 180)) + randnum(),\n",
+ " random.choice(range(-89, 90)) + randnum()\n",
+ " ) for _ in range(n_samples)\n",
+ "]\n",
"kd_res = [kt.query(r) for r in test_records]\n",
"kd_recs = [_[0][0] for _ in kd_res]\n",
"kd_dists = [_[1] for _ in kd_res]\n",
- "tr_recs = [records[np.argmin([r.distance(p) for p in records])] for r in test_records]\n",
+ "tr_recs = [\n",
+ " records[np.argmin([r.distance(p) for p in records])]\n",
+ " for r in test_records\n",
+ "]\n",
"tr_dists = [min([r.distance(p) for p in records]) for r in test_records]\n",
- "assert all([abs(k - t) < tol for k, t in zip(kd_dists, tr_dists)]), \"NOT MATCHING?\""
+ "\n",
+ "if not all([abs(k - t) < tol for k, t in zip(kd_dists, tr_dists)]):\n",
+ " raise ValueError(\"NOT MATCHING?\")"
]
},
{
"cell_type": "code",
- "execution_count": 10,
+ "execution_count": 13,
"id": "a6aa6926-7fd5-4fff-bd20-7bc0305b948d",
"metadata": {},
"outputs": [
@@ -214,7 +480,7 @@
"└──────────┴──────────┴─────────┴────────┴────────┴─────────┴────────┴────────┘"
]
},
- "execution_count": 10,
+ "execution_count": 13,
"metadata": {},
"output_type": "execute_result"
}
@@ -230,14 +496,14 @@
"tr_lats = [r.lat for r in tr_recs]\n",
"\n",
"df = pl.DataFrame({\n",
- " \"test_lon\": test_lons, \n",
+ " \"test_lon\": test_lons,\n",
" \"test_lat\": test_lats,\n",
" \"kd_dist\": kd_dists,\n",
" \"kd_lon\": kd_lons,\n",
" \"kd_lat\": kd_lats,\n",
" \"tr_dist\": tr_dists,\n",
" \"tr_lon\": tr_lons,\n",
- " \"tr_lat\": tr_lats, \n",
+ " \"tr_lat\": tr_lats,\n",
"}).filter((pl.col(\"kd_dist\") - pl.col(\"tr_dist\")).abs().ge(tol))\n",
"df"
]
@@ -245,7 +511,7 @@
],
"metadata": {
"kernelspec": {
- "display_name": "GeoSpatialTools",
+ "display_name": "geospatialtools",
"language": "python",
"name": "geospatialtools"
},
@@ -259,7 +525,7 @@
"name": "python",
"nbconvert_exporter": "python",
"pygments_lexer": "ipython3",
- "version": "3.12.7"
+ "version": "3.11.11"
}
},
"nbformat": 4,
diff --git a/notebooks/octtree.ipynb b/notebooks/octtree.ipynb
index 1259795099391eeda11058b5932dd0c5d975f8f6..474cb2f3bab9c0b3fae610ea31945fda434807fb 100644
--- a/notebooks/octtree.ipynb
+++ b/notebooks/octtree.ipynb
@@ -7,7 +7,9 @@
"source": [
"## OctTree!\n",
"\n",
- "Testing the time to look-up nearby records with the PyCOADS OctTree implementation."
+ "Testing the time to look-up nearby records with the `OctTree` implementation.\n",
+ "\n",
+ "The `OctTree` is used to find records within a spatio-temporal range of a given point, or within a box defined by lon, lat, & time bounds."
]
},
{
@@ -18,11 +20,9 @@
"outputs": [],
"source": [
"import os\n",
- "import gzip\n",
"os.environ[\"POLARS_MAX_THREADS\"] = \"4\"\n",
"\n",
"from datetime import datetime, timedelta\n",
- "from random import choice\n",
"from string import ascii_letters, digits\n",
"import random\n",
"import inspect\n",
@@ -30,295 +30,28 @@
"import polars as pl\n",
"import numpy as np\n",
"\n",
- "from GeoSpatialTools.octtree import OctTree, SpaceTimeRecord as Record, SpaceTimeRectangle as Rectangle"
- ]
- },
- {
- "cell_type": "raw",
- "id": "99295bad-0db3-444b-8d38-acc7875cc0f0",
- "metadata": {},
- "source": [
- "## Generate Data\n",
- "\n",
- "16,000 rows of data"
- ]
- },
- {
- "cell_type": "code",
- "execution_count": 2,
- "id": "d8f1e5e1-513c-4bdf-a9f9-cef9562a7cb7",
- "metadata": {},
- "outputs": [],
- "source": [
- "def generate_uid(n: int) -> str:\n",
- " chars = ascii_letters + digits\n",
- " return \"\".join(random.choice(chars) for _ in range(n))"
- ]
- },
- {
- "cell_type": "code",
- "execution_count": 3,
- "id": "986d9cc5-e610-449a-9ee7-e281b7558ca9",
- "metadata": {},
- "outputs": [],
- "source": [
- "N = 16_000\n",
- "lons = pl.int_range(-180, 180, eager=True)\n",
- "lats = pl.int_range(-90, 90, eager=True)\n",
- "dates = pl.datetime_range(datetime(1900, 1, 1, 0), datetime(1900, 1, 31, 23), interval=\"1h\", eager=True)\n",
- "\n",
- "lons_use = lons.sample(N, with_replacement=True).alias(\"lon\")\n",
- "lats_use = lats.sample(N, with_replacement=True).alias(\"lat\")\n",
- "dates_use = dates.sample(N, with_replacement=True).alias(\"datetime\")\n",
- "uids = pl.Series(\"uid\", [generate_uid(8) for _ in range(N)])\n",
- "\n",
- "df = pl.DataFrame([lons_use, lats_use, dates_use, uids]).unique()"
- ]
- },
- {
- "cell_type": "markdown",
- "id": "237096f1-093e-49f0-9a9a-2bec5231726f",
- "metadata": {},
- "source": [
- "## Add extra rows\n",
- "\n",
- "For testing larger datasets. Uncomment to use."
- ]
- },
- {
- "cell_type": "code",
- "execution_count": 4,
- "id": "0b8fd425-8a90-4f76-91b7-60df48aa98e4",
- "metadata": {},
- "outputs": [],
- "source": [
- "# _df = df.clone()\n",
- "# for i in range(100):\n",
- "# df2 = pl.DataFrame([\n",
- "# _df[\"lon\"].shuffle(),\n",
- "# _df[\"lat\"].shuffle(),\n",
- "# _df[\"datetime\"].shuffle(),\n",
- "# _df[\"uid\"].shuffle(),\n",
- "# ]).with_columns(pl.concat_str([pl.col(\"uid\"), pl.lit(f\"{i:03d}\")]).alias(\"uid\"))\n",
- "# df = df.vstack(df2)\n",
- "# df.shape\n",
- "# df"
- ]
- },
- {
- "cell_type": "markdown",
- "id": "c7bd16e0-96a6-426b-b00a-7c3b8a2aaddd",
- "metadata": {},
- "source": [
- "## Intialise the OctTree Object"
- ]
- },
- {
- "cell_type": "code",
- "execution_count": 5,
- "id": "af06a976-ff52-49e0-a886-91bcbe540ffe",
- "metadata": {},
- "outputs": [],
- "source": [
- "otree = OctTree(Rectangle(-180, 180, -90, 90, datetime(1900, 1, 1, 0), datetime(1900, 1, 31, 23)), capacity = 10, max_depth = 25)"
- ]
- },
- {
- "cell_type": "code",
- "execution_count": 6,
- "id": "2ba99b37-787c-4862-8075-a7596208c60e",
- "metadata": {},
- "outputs": [
- {
- "name": "stdout",
- "output_type": "stream",
- "text": [
- "CPU times: user 106 ms, sys: 3.98 ms, total: 110 ms\n",
- "Wall time: 109 ms\n"
- ]
- }
- ],
- "source": [
- "%%time\n",
- "for r in df.rows():\n",
- " otree.insert(Record(*r))"
- ]
- },
- {
- "cell_type": "code",
- "execution_count": 7,
- "id": "59d38446-f7d2-4eec-bba3-c39bd7279623",
- "metadata": {
- "scrolled": true
- },
- "outputs": [
- {
- "name": "stdout",
- "output_type": "stream",
- "text": [
- "OctTree:\n",
- "- boundary: SpaceTimeRectangle(west=-180, east=180, south=-90, north=90, start=datetime.datetime(1900, 1, 1, 0, 0), end=datetime.datetime(1900, 1, 31, 23, 0))\n",
- "- capacity: 10\n",
- "- depth: 0\n",
- "- max_depth: 25\n",
- "- contents:\n",
- "- number of elements: 10\n",
- " * SpaceTimeRecord(x = 92, y = 15, datetime = 1900-01-17 08:00:00, uid = HRF401hH)\n",
- " * SpaceTimeRecord(x = -35, y = 37, datetime = 1900-01-04 08:00:00, uid = CXZaSOdh)\n",
- " * SpaceTimeRecord(x = 84, y = -7, datetime = 1900-01-07 16:00:00, uid = 2aEjxGwG)\n",
- " * SpaceTimeRecord(x = 68, y = 73, datetime = 1900-01-18 17:00:00, uid = Ah7lanWB)\n",
- " * SpaceTimeRecord(x = -179, y = 40, datetime = 1900-01-01 11:00:00, uid = HGxSJzf4)\n",
- " * SpaceTimeRecord(x = -73, y = 23, datetime = 1900-01-09 12:00:00, uid = qHQ8opO9)\n",
- " * SpaceTimeRecord(x = 117, y = -23, datetime = 1900-01-31 06:00:00, uid = ctvs56Fq)\n",
- " * SpaceTimeRecord(x = 109, y = 55, datetime = 1900-01-13 14:00:00, uid = C2xXIglD)\n",
- " * SpaceTimeRecord(x = 104, y = -10, datetime = 1900-01-06 16:00:00, uid = WEpQKIOV)\n",
- " * SpaceTimeRecord(x = 45, y = -71, datetime = 1900-01-29 00:00:00, uid = 7r1UeXRi)\n",
- "- with children:\n",
- " OctTree:\n",
- " - boundary: SpaceTimeRectangle(west=-180, east=0.0, south=0.0, north=90, start=datetime.datetime(1900, 1, 1, 0, 0), end=datetime.datetime(1900, 1, 16, 11, 30))\n",
- " - capacity: 10\n",
- " - depth: 1\n",
- " - max_depth: 25\n",
- " - contents:\n",
- " - number of elements: 10\n",
- " * SpaceTimeRecord(x = -84, y = 38, datetime = 1900-01-15 10:00:00, uid = 63mpq3Kx)\n",
- " * SpaceTimeRecord(x = -78, y = 60, datetime = 1900-01-10 01:00:00, uid = vZ8HLu5t)\n",
- " * SpaceTimeRecord(x = -89, y = 24, datetime = 1900-01-12 17:00:00, uid = gn2o9tYQ)\n",
- " * SpaceTimeRecord(x = -149, y = 7, datetime = 1900-01-08 11:00:00, uid = 2ODnGJO6)\n",
- " * SpaceTimeRecord(x = -37, y = 54, datetime = 1900-01-12 13:00:00, uid = 11cApOwm)\n",
- " * SpaceTimeRecord(x = -34, y = 88, datetime = 1900-01-03 05:00:00, uid = 8SN6zPWh)\n",
- " * SpaceTimeRecord(x = -36, y = 13, datetime = 1900-01-14 13:00:00, uid = ijfjmp8E)\n",
- " * SpaceTimeRecord(x = -168, y = 62, datetime = 1900-01-03 09:00:00, uid = Cc4m1azR)\n",
- " * SpaceTimeRecord(x = -76, y = 67, datetime = 1900-01-06 04:00:00, uid = 4WeWpZUz)\n",
- " * SpaceTimeRecord(x = -156, y = 39, datetime = 1900-01-13 10:00:00, uid = dZXAMaXq)\n",
- " - with children:\n",
- " OctTree:\n",
- " - boundary: SpaceTimeRectangle(west=-180, east=-90.0, south=45.0, north=90, start=datetime.datetime(1900, 1, 1, 0, 0), end=datetime.datetime(1900, 1, 8, 17, 45))\n",
- " - capacity: 10\n",
- " - depth: 2\n",
- " - max_depth: 25\n",
- " - contents:\n",
- " - number of elements: 10\n",
- " * SpaceTimeRecord(x = -141, y = 79, datetime = 1900-01-03 05:00:00, uid = mN1Mg7Vn)\n",
- " * SpaceTimeRecord(x = -172, y = 80, datetime = 1900-01-01 14:00:00, uid = NBBZ3bCW)\n",
- " * SpaceTimeRecord(x = -93, y = 53, datetime = 1900-01-06 07:00:00, uid = jX8HZPJT)\n",
- " * SpaceTimeRecord(x = -168, y = 82, datetime = 1900-01-03 08:00:00, uid = dlxpN1Ew)\n",
- " * SpaceTimeRecord(x = -111, y = 83, datetime = 1900-01-02 12:00:00, uid = GXLopHH0)\n",
- " * SpaceTimeRecord(x = -178, y = 61, datetime = 1900-01-02 00:00:00, uid = 0ut6CLe5)\n",
- " * SpaceTimeRecord(x = -148, y = 74, datetime = 1900-01-07 23:00:00, uid = xUySW1tx)\n",
- " * SpaceTimeRecord(x = -174, y = 63, datetime = 1900-01-06 22:00:00, uid = 8sI94Lt6)\n",
- " * SpaceTimeRecord(x = -114, y = 84, datetime = 1900-01-05 15:00:00, uid = OoY9mEkQ)\n",
- " * SpaceTimeRecord(x = -102, y = 82, datetime = 1900-01-02 15:00:00, uid = bd4sLang)\n",
- " - with children:\n",
- " OctTree:\n",
- " - boundary: SpaceTimeRectangle(west=-180, east=-135.0, south=67.5, north=90, start=datetime.datetime(1900, 1, 1, 0, 0), end=datetime.datetime(1900, 1, 4, 20, 52, 30))\n",
- " - capacity: 10\n",
- " - depth: 3\n",
- " - max_depth: 25\n",
- " - contents:\n",
- " - number of elements: 10\n",
- " * SpaceTimeRecord(x = -148, y = 79, datetime = 1900-01-03 21:00:00, uid = kNWm70rm)\n",
- " * SpaceTimeRecord(x = -157, y = 80, datetime = 1900-01-03 05:00:00, uid = 471X27tA)\n",
- " * SpaceTimeRecord(x = -152, y = 85, datetime = 1900-01-03 01:00:00, uid = cjTyQn7E)\n",
- " * SpaceTimeRecord(x = -154, y = 88, datetime = 1900-01-03 15:00:00, uid = JTnjCJZN)\n",
- " * SpaceTimeRecord(x = -139, y = 83, datetime = 1900-01-01 21:00:00, uid = kZ28j8I5)\n",
- " * SpaceTimeRecord(x = -161, y = 73, datetime = 1900-01-03 02:00:00, uid = wsHJBLLC)\n",
- " * SpaceTimeRecord(x = -140, y = 71, datetime = 1900-01-02 07:00:00, uid = 4bTg1N2k)\n",
- " * SpaceTimeRecord(x = -141, y = 74, datetime = 1900-01-04 09:00:00, uid = I6M8kuue)\n",
- " * SpaceTimeRecord(x = -144, y = 72, datetime = 1900-01-04 17:00:00, uid = 0fPvYOC9)\n",
- " * SpaceTimeRecord(x = -157, y = 78, datetime = 1900-01-03 16:00:00, uid = yAL3OeaK)\n",
- " - with children:\n",
- " OctTree:\n",
- " - boundary: SpaceTimeRectangle(west=-180, east=-157.5, south=78.75, north=90, start=datetime.datetime(1900, 1, 1, 0, 0), end=datetime.datetime(1900, 1, 2, 22, 26, 15))\n",
- " - capacity: 10\n",
- " - depth: 4\n",
- " - max_depth: 25\n",
- " - contents:\n",
- " - number of elements: 4\n",
- " * SpaceTimeRecord(x = -180, y = 88, datetime = 1900-01-02 12:00:00, uid = CXeAd3y4)\n",
- " * SpaceTimeRecord(x = -180, y = 87, datetime = 1900-01-01 16:00:00, uid = TB2xKFgK)\n",
- " * SpaceTimeRecord(x = -171, y = 79, datetime = 1900-01-02 04:00:00, uid = pIU8qvxT)\n",
- " * SpaceTimeRecord(x = -168, y = 85, datetime = 1900-01-01 22:00:00, uid = 7zL4gz8K)\n",
- " OctTree:\n",
- " - boundary: SpaceTimeRectangle(west=-157.5, east=-135.0, south=78.75, north=90, start=datetime.datetime(1900, 1, 1, 0, 0), end=datetime.datetime(1900, 1, 2, 22, 26, 15))\n",
- " - capacity: 10\n",
- " - depth: 4\n",
- " - max_depth: 25\n",
- " - contents:\n",
- " - number of elements: 2\n",
- " * SpaceTimeRecord(x = -149, y = 82, datetime = 1900-01-01 20:00:00, uid = xTYMs6Xp)\n",
- " * SpaceTimeRecord(x = -154, y = 84, datetime = 1900-01-02 21:00:00, uid = JSEaGBsn)\n",
- " OctTree:\n",
- " - boundary: SpaceTimeRectangle(west=-180, east=-157.5, south=67.5, north=78.75, start=datetime.datetime(1900, 1, 1, 0, 0), end=datetime.datetime(1900, 1, 2, 22, 26, 15))\n",
- " - capacity: 10\n",
- " - depth: 4\n",
- " - max_depth: 25\n",
- " - contents:\n",
- " - number of elements: 3\n",
- " * SpaceTimeRecord(x = -173, y = 75, datetime = 1900-01-01 06:00:00, uid = M4N3amQ3)\n"
- ]
- }
- ],
- "source": [
- "s = str(otree)\n",
- "print(\"\\n\".join(s.split(\"\\n\")[:100]))"
+ "from GeoSpatialTools.octtree import (\n",
+ " OctTree,\n",
+ " SpaceTimeRecord as Record,\n",
+ " SpaceTimeRectangle as Rectangle\n",
+ ")"
]
},
{
"cell_type": "markdown",
- "id": "6b02c2ea-6566-47c2-97e0-43d8b18e0713",
+ "id": "6b0e8015-b958-4be7-9e63-9e21f081011b",
"metadata": {},
"source": [
- "## Time Execution\n",
+ "## Set-up functions\n",
"\n",
- "Testing the identification of nearby points against the original full search"
- ]
- },
- {
- "cell_type": "code",
- "execution_count": 8,
- "id": "094b588c-e938-4838-9719-1defdfff74fa",
- "metadata": {},
- "outputs": [],
- "source": [
- "dts = pl.datetime_range(datetime(1900, 1, 1), datetime(1900, 2, 1), interval=\"1h\", eager=True, closed=\"left\")\n",
- "N = dts.len()\n",
- "lons = 180 - 360 * np.random.rand(N)\n",
- "lats = 90 - 180 * np.random.rand(N)\n",
- "test_df = pl.DataFrame({\"lon\": lons, \"lat\": lats, \"datetime\": dts})\n",
- "test_recs = [Record(*r) for r in test_df.rows()]\n",
- "dt = timedelta(days = 1)\n",
- "dist = 350"
- ]
- },
- {
- "cell_type": "code",
- "execution_count": 9,
- "id": "66a48b86-d449-45d2-9837-2b3e07f5563d",
- "metadata": {},
- "outputs": [
- {
- "name": "stdout",
- "output_type": "stream",
- "text": [
- "207 μs ± 6.25 μs per loop (mean ± std. dev. of 7 runs, 1,000 loops each)\n"
- ]
- }
- ],
- "source": [
- "%%timeit\n",
- "otree.nearby_points(random.choice(test_recs), dist=dist, t_dist=dt)"
+ "For comparisons using brute-force appraoch"
]
},
{
"cell_type": "code",
- "execution_count": 10,
+ "execution_count": 2,
"id": "972d4a16-39fd-4f80-8592-1c5d5cabf5be",
- "metadata": {
- "jupyter": {
- "source_hidden": true
- }
- },
+ "metadata": {},
"outputs": [],
"source": [
"def check_cols(\n",
@@ -353,7 +86,6 @@
" err_str = f\"({calling_func}) - {var_name} missing required columns. \"\n",
" err_str += f'Require: {\", \".join(cols)}. '\n",
" err_str += f'Missing: {\", \".join(missing)}.'\n",
- " logging.error(err_str)\n",
" raise ValueError(err_str)\n",
"\n",
" return\n",
@@ -361,24 +93,15 @@
"\n",
"def haversine_df(\n",
" df: pl.DataFrame | pl.LazyFrame,\n",
- " date_var: str = \"datetime\",\n",
+ " lon: float,\n",
+ " lat: float,\n",
" R: float = 6371,\n",
- " reverse: bool = False,\n",
- " out_colname: str = \"dist\",\n",
" lon_col: str = \"lon\",\n",
" lat_col: str = \"lat\",\n",
- " lon2_col: str | None = None,\n",
- " lat2_col: str | None = None,\n",
- " sorted: bool = False,\n",
- " rev_prefix: str = \"rev_\",\n",
") -> pl.DataFrame | pl.LazyFrame:\n",
" \"\"\"\n",
- " Compute haversine distance on earth surface between lon-lat positions.\n",
- "\n",
- " If only 'lon_col' and 'lat_col' are specified then this computes the\n",
- " distance between consecutive points. If a second set of positions is\n",
- " included via the optional 'lon2_col' and 'lat2_col' arguments then the\n",
- " distances between the columns are computed.\n",
+ " Compute haversine distance on earth surface between lon-lat positions\n",
+ " in a polars DataFrame and a lon-lat position.\n",
"\n",
" Parameters\n",
" ----------\n",
@@ -387,27 +110,16 @@
" * lon_col\n",
" * lat_col\n",
" * date_var\n",
- " date_var : str\n",
- " Name of the datetime column on which to sort the positions\n",
+ " lon : float\n",
+ " The longitude of the position.\n",
+ " lat : float\n",
+ " The latitude of the position.\n",
" R : float\n",
" Radius of earth in km\n",
- " reverse : bool\n",
- " Compute distances in reverse\n",
- " out_colname : str\n",
- " Name of the output column to store distances. Prefixed with 'rev_' if\n",
- " reverse is True\n",
" lon_col : str\n",
" Name of the longitude column\n",
" lat_col : str\n",
" Name of the latitude column\n",
- " lon2_col : str\n",
- " Name of the 2nd longitude column if present\n",
- " lat2_col : str\n",
- " Name of the 2nd latitude column if present\n",
- " sorted : bool\n",
- " Compute distances assuming that the frame is already sorted\n",
- " rev_prefix : str\n",
- " Prefix to use for colnames if reverse is True\n",
"\n",
" Returns\n",
" -------\n",
@@ -417,61 +129,14 @@
" \"\"\"\n",
" required_cols = [lon_col, lat_col]\n",
"\n",
- " if lon2_col is not None and lat2_col is not None:\n",
- " required_cols += [lon2_col, lat2_col]\n",
- " check_cols(df, required_cols, \"df\")\n",
- " return (\n",
- " df.with_columns(\n",
- " [\n",
- " pl.col(lat_col).radians().alias(\"_lat0\"),\n",
- " pl.col(lat2_col).radians().alias(\"_lat1\"),\n",
- " (pl.col(lon_col) - pl.col(lon2_col))\n",
- " .radians()\n",
- " .alias(\"_dlon\"),\n",
- " (pl.col(lat_col) - pl.col(lat2_col))\n",
- " .radians()\n",
- " .alias(\"_dlat\"),\n",
- " ]\n",
- " )\n",
- " .with_columns(\n",
- " (\n",
- " (pl.col(\"_dlat\") / 2).sin().pow(2)\n",
- " + pl.col(\"_lat0\").cos()\n",
- " * pl.col(\"_lat1\").cos()\n",
- " * (pl.col(\"_dlon\") / 2).sin().pow(2)\n",
- " ).alias(\"_a\")\n",
- " )\n",
- " .with_columns(\n",
- " (2 * R * (pl.col(\"_a\").sqrt().arcsin()))\n",
- " .round(2)\n",
- " .alias(out_colname)\n",
- " )\n",
- " .drop([\"_lat0\", \"_lat1\", \"_dlon\", \"_dlat\", \"_a\"])\n",
- " )\n",
- "\n",
- " if lon2_col is not None or lat2_col is not None:\n",
- " logging.warning(\n",
- " \"(haversine_df) 2nd position incorrectly specified. \"\n",
- " + \"Calculating consecutive distances.\"\n",
- " )\n",
- "\n",
- " required_cols += [date_var]\n",
" check_cols(df, required_cols, \"df\")\n",
- " if reverse:\n",
- " out_colname = rev_prefix + out_colname\n",
- " if not sorted:\n",
- " df = df.sort(date_var, descending=reverse)\n",
" return (\n",
" df.with_columns(\n",
" [\n",
" pl.col(lat_col).radians().alias(\"_lat0\"),\n",
- " pl.col(lat_col).shift(n=-1).radians().alias(\"_lat1\"),\n",
- " (pl.col(lon_col).shift(n=-1) - pl.col(lon_col))\n",
- " .radians()\n",
- " .alias(\"_dlon\"),\n",
- " (pl.col(lat_col).shift(n=-1) - pl.col(lat_col))\n",
- " .radians()\n",
- " .alias(\"_dlat\"),\n",
+ " pl.lit(lat).radians().alias(\"_lat1\"),\n",
+ " (pl.col(lon_col) - lon).radians().alias(\"_dlon\"),\n",
+ " (pl.col(lat_col) - lat).radians().alias(\"_dlat\"),\n",
" ]\n",
" )\n",
" .with_columns(\n",
@@ -485,15 +150,17 @@
" .with_columns(\n",
" (2 * R * (pl.col(\"_a\").sqrt().arcsin()))\n",
" .round(2)\n",
- " .fill_null(strategy=\"forward\")\n",
- " .alias(out_colname)\n",
+ " .alias(\"_dist\")\n",
" )\n",
" .drop([\"_lat0\", \"_lat1\", \"_dlon\", \"_dlat\", \"_a\"])\n",
" )\n",
"\n",
+ "\n",
"def intersect(a, b) -> set:\n",
+ " \"\"\"Intersection of a and b, items in both a and b\"\"\"\n",
" return set(a) & set(b)\n",
"\n",
+ "\n",
"def nearby_ships(\n",
" lon: float,\n",
" lat: float,\n",
@@ -573,26 +240,200 @@
" )\n",
"\n",
" return (\n",
- " pool.with_columns(\n",
- " [pl.lit(lon).alias(\"_lon\"), pl.lit(lat).alias(\"_lat\")]\n",
- " )\n",
+ " pool\n",
" .pipe(\n",
" haversine_df,\n",
+ " lon=lon,\n",
+ " lat=lat,\n",
" lon_col=lon_col,\n",
" lat_col=lat_col,\n",
- " out_colname=\"_dist\",\n",
- " lon2_col=\"_lon\",\n",
- " lat2_col=\"_lat\",\n",
" )\n",
" .filter(pl.col(\"_dist\").le(max_dist))\n",
- " .drop([\"_dist\", \"_lon\", \"_lat\"])\n",
+ " .drop([\"_dist\"])\n",
" )\n"
]
},
+ {
+ "cell_type": "markdown",
+ "id": "08cba819-1ebe-48b3-85c7-3fd7469399f8",
+ "metadata": {},
+ "source": [
+ "## Generate Data\n",
+ "\n",
+ "16,000 rows of data"
+ ]
+ },
{
"cell_type": "code",
- "execution_count": 11,
- "id": "8b9279ed-6f89-4423-8833-acd0b365eb7b",
+ "execution_count": 3,
+ "id": "d8f1e5e1-513c-4bdf-a9f9-cef9562a7cb7",
+ "metadata": {},
+ "outputs": [],
+ "source": [
+ "def generate_uid(n: int) -> str:\n",
+ " \"\"\"Generates a psuedo uid by randomly selecting from characters\"\"\"\n",
+ " chars = ascii_letters + digits\n",
+ " return \"\".join(random.choice(chars) for _ in range(n))"
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": 4,
+ "id": "986d9cc5-e610-449a-9ee7-e281b7558ca9",
+ "metadata": {},
+ "outputs": [],
+ "source": [
+ "N = 16_000\n",
+ "lons = pl.int_range(-180, 180, eager=True)\n",
+ "lats = pl.int_range(-90, 90, eager=True)\n",
+ "dates = pl.datetime_range(\n",
+ " datetime(1900, 1, 1, 0),\n",
+ " datetime(1900, 1, 31, 23),\n",
+ " interval=\"1h\",\n",
+ " eager=True,\n",
+ ")\n",
+ "\n",
+ "lons_use = lons.sample(N, with_replacement=True).alias(\"lon\")\n",
+ "lats_use = lats.sample(N, with_replacement=True).alias(\"lat\")\n",
+ "dates_use = dates.sample(N, with_replacement=True).alias(\"datetime\")\n",
+ "uids = pl.Series(\"uid\", [generate_uid(8) for _ in range(N)])\n",
+ "\n",
+ "df = pl.DataFrame([lons_use, lats_use, dates_use, uids]).unique()"
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "id": "237096f1-093e-49f0-9a9a-2bec5231726f",
+ "metadata": {},
+ "source": [
+ "## Add extra rows\n",
+ "\n",
+ "For testing larger datasets. Uncomment to use."
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": 5,
+ "id": "0b8fd425-8a90-4f76-91b7-60df48aa98e4",
+ "metadata": {},
+ "outputs": [
+ {
+ "data": {
+ "text/html": [
+ "<div><style>\n",
+ ".dataframe > thead > tr,\n",
+ ".dataframe > tbody > tr {\n",
+ " text-align: right;\n",
+ " white-space: pre-wrap;\n",
+ "}\n",
+ "</style>\n",
+ "<small>shape: (1_616_000, 4)</small><table border=\"1\" class=\"dataframe\"><thead><tr><th>lon</th><th>lat</th><th>datetime</th><th>uid</th></tr><tr><td>i64</td><td>i64</td><td>datetime[μs]</td><td>str</td></tr></thead><tbody><tr><td>-161</td><td>6</td><td>1900-01-22 17:00:00</td><td>"qCulqvN6"</td></tr><tr><td>-1</td><td>25</td><td>1900-01-23 17:00:00</td><td>"krL2tTTH"</td></tr><tr><td>146</td><td>-20</td><td>1900-01-08 22:00:00</td><td>"QCASMObF"</td></tr><tr><td>-16</td><td>-38</td><td>1900-01-05 05:00:00</td><td>"Wh9pptMZ"</td></tr><tr><td>-127</td><td>-33</td><td>1900-01-10 20:00:00</td><td>"PPIxvkbU"</td></tr><tr><td>…</td><td>…</td><td>…</td><td>…</td></tr><tr><td>62</td><td>0</td><td>1900-01-22 11:00:00</td><td>"6PxQzuHv099"</td></tr><tr><td>90</td><td>-34</td><td>1900-01-16 00:00:00</td><td>"pyjCOpgo099"</td></tr><tr><td>-132</td><td>-20</td><td>1900-01-23 05:00:00</td><td>"vZLzl0aX099"</td></tr><tr><td>-104</td><td>39</td><td>1900-01-06 15:00:00</td><td>"8kavFVpP099"</td></tr><tr><td>-131</td><td>-81</td><td>1900-01-24 19:00:00</td><td>"LFTv3XZ1099"</td></tr></tbody></table></div>"
+ ],
+ "text/plain": [
+ "shape: (1_616_000, 4)\n",
+ "┌──────┬─────┬─────────────────────┬─────────────â”\n",
+ "│ lon ┆ lat ┆ datetime ┆ uid │\n",
+ "│ --- ┆ --- ┆ --- ┆ --- │\n",
+ "│ i64 ┆ i64 ┆ datetime[μs] ┆ str │\n",
+ "â•žâ•â•â•â•â•â•â•ªâ•â•â•â•â•â•ªâ•â•â•â•â•â•â•â•â•â•â•â•â•â•â•â•â•â•â•â•â•â•ªâ•â•â•â•â•â•â•â•â•â•â•â•â•â•¡\n",
+ "│ -161 ┆ 6 ┆ 1900-01-22 17:00:00 ┆ qCulqvN6 │\n",
+ "│ -1 ┆ 25 ┆ 1900-01-23 17:00:00 ┆ krL2tTTH │\n",
+ "│ 146 ┆ -20 ┆ 1900-01-08 22:00:00 ┆ QCASMObF │\n",
+ "│ -16 ┆ -38 ┆ 1900-01-05 05:00:00 ┆ Wh9pptMZ │\n",
+ "│ -127 ┆ -33 ┆ 1900-01-10 20:00:00 ┆ PPIxvkbU │\n",
+ "│ … ┆ … ┆ … ┆ … │\n",
+ "│ 62 ┆ 0 ┆ 1900-01-22 11:00:00 ┆ 6PxQzuHv099 │\n",
+ "│ 90 ┆ -34 ┆ 1900-01-16 00:00:00 ┆ pyjCOpgo099 │\n",
+ "│ -132 ┆ -20 ┆ 1900-01-23 05:00:00 ┆ vZLzl0aX099 │\n",
+ "│ -104 ┆ 39 ┆ 1900-01-06 15:00:00 ┆ 8kavFVpP099 │\n",
+ "│ -131 ┆ -81 ┆ 1900-01-24 19:00:00 ┆ LFTv3XZ1099 │\n",
+ "└──────┴─────┴─────────────────────┴─────────────┘"
+ ]
+ },
+ "execution_count": 5,
+ "metadata": {},
+ "output_type": "execute_result"
+ }
+ ],
+ "source": [
+ "_df = df.clone()\n",
+ "for i in range(100):\n",
+ " df2 = pl.DataFrame([\n",
+ " _df[\"lon\"].shuffle(),\n",
+ " _df[\"lat\"].shuffle(),\n",
+ " _df[\"datetime\"].shuffle(),\n",
+ " _df[\"uid\"].shuffle(),\n",
+ " ]).with_columns(\n",
+ " pl.concat_str([pl.col(\"uid\"), pl.lit(f\"{i:03d}\")]).alias(\"uid\")\n",
+ " )\n",
+ " df = df.vstack(df2)\n",
+ "df.shape\n",
+ "df"
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "id": "c7bd16e0-96a6-426b-b00a-7c3b8a2aaddd",
+ "metadata": {},
+ "source": [
+ "## Intialise the OctTree Object\n",
+ "\n",
+ "There is an overhead in constructing the `OctTree` class. The performance benefits appear if multiple neighbourhood searches are performed."
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": 6,
+ "id": "af06a976-ff52-49e0-a886-91bcbe540ffe",
+ "metadata": {},
+ "outputs": [],
+ "source": [
+ "bounds = Rectangle(\n",
+ " -180,\n",
+ " 180,\n",
+ " -90,\n",
+ " 90,\n",
+ " datetime(1900, 1, 1, 0),\n",
+ " datetime(1900, 1, 31, 23),\n",
+ ")\n",
+ "otree = OctTree(bounds, capacity=10, max_depth=25)"
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": 7,
+ "id": "2ba99b37-787c-4862-8075-a7596208c60e",
+ "metadata": {},
+ "outputs": [
+ {
+ "name": "stdout",
+ "output_type": "stream",
+ "text": [
+ "CPU times: user 16.3 s, sys: 163 ms, total: 16.4 s\n",
+ "Wall time: 16.5 s\n"
+ ]
+ }
+ ],
+ "source": [
+ "%%time\n",
+ "for r in df.rows():\n",
+ " otree.insert(Record(*r))"
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "id": "94be9d8a-02fc-49f2-98c9-0bcf250b1d10",
+ "metadata": {},
+ "source": [
+ "### View the `OctTree`\n",
+ "\n",
+ "It is a nested object, with child `OctTree` objects in each octant of the space-time domain."
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": 8,
+ "id": "59d38446-f7d2-4eec-bba3-c39bd7279623",
"metadata": {
"scrolled": true
},
@@ -601,47 +442,219 @@
"name": "stdout",
"output_type": "stream",
"text": [
- "5.36 ms ± 164 μs per loop (mean ± std. dev. of 7 runs, 100 loops each)\n"
+ "OctTree:\n",
+ "- boundary: SpaceTimeRectangle(west=-180, east=180, south=-90, north=90, start=datetime.datetime(1900, 1, 1, 0, 0), end=datetime.datetime(1900, 1, 31, 23, 0))\n",
+ "- capacity: 10\n",
+ "- depth: 0\n",
+ "- max_depth: 25\n",
+ "- contents:\n",
+ "- number of elements: 10\n",
+ " * SpaceTimeRecord(x = -161, y = 6, datetime = 1900-01-22 17:00:00, uid = qCulqvN6)\n",
+ " * SpaceTimeRecord(x = -1, y = 25, datetime = 1900-01-23 17:00:00, uid = krL2tTTH)\n",
+ " * SpaceTimeRecord(x = 146, y = -20, datetime = 1900-01-08 22:00:00, uid = QCASMObF)\n",
+ " * SpaceTimeRecord(x = -16, y = -38, datetime = 1900-01-05 05:00:00, uid = Wh9pptMZ)\n",
+ " * SpaceTimeRecord(x = -127, y = -33, datetime = 1900-01-10 20:00:00, uid = PPIxvkbU)\n",
+ " * SpaceTimeRecord(x = 88, y = 37, datetime = 1900-01-18 12:00:00, uid = gYAwqD2R)\n",
+ " * SpaceTimeRecord(x = -122, y = 57, datetime = 1900-01-14 13:00:00, uid = L77bWRL1)\n",
+ " * SpaceTimeRecord(x = -179, y = 23, datetime = 1900-01-29 23:00:00, uid = 3jSwN6aK)\n",
+ " * SpaceTimeRecord(x = -156, y = 79, datetime = 1900-01-25 16:00:00, uid = OYEzYral)\n",
+ " * SpaceTimeRecord(x = 140, y = 15, datetime = 1900-01-07 20:00:00, uid = dNqilTiD)\n",
+ "- with children:\n",
+ " OctTree:\n",
+ " - boundary: SpaceTimeRectangle(west=-180, east=0.0, south=0.0, north=90, start=datetime.datetime(1900, 1, 1, 0, 0), end=datetime.datetime(1900, 1, 16, 11, 30))\n",
+ " - capacity: 10\n",
+ " - depth: 1\n",
+ " - max_depth: 25\n",
+ " - contents:\n",
+ " - number of elements: 10\n",
+ " * SpaceTimeRecord(x = -107, y = 89, datetime = 1900-01-10 01:00:00, uid = dB4jDgBL)\n",
+ " * SpaceTimeRecord(x = -132, y = 50, datetime = 1900-01-11 18:00:00, uid = ZOzYoDbB)\n",
+ " * SpaceTimeRecord(x = -28, y = 5, datetime = 1900-01-13 17:00:00, uid = YwB5kPdG)\n",
+ " * SpaceTimeRecord(x = -153, y = 25, datetime = 1900-01-10 22:00:00, uid = vzNI6J3z)\n",
+ " * SpaceTimeRecord(x = -45, y = 12, datetime = 1900-01-13 18:00:00, uid = kwHmr9mE)\n",
+ " * SpaceTimeRecord(x = -31, y = 16, datetime = 1900-01-06 17:00:00, uid = h3JQR5Ab)\n",
+ " * SpaceTimeRecord(x = -153, y = 25, datetime = 1900-01-14 03:00:00, uid = ZgZwvzHY)\n",
+ " * SpaceTimeRecord(x = -142, y = 43, datetime = 1900-01-15 14:00:00, uid = jd0JycvC)\n",
+ " * SpaceTimeRecord(x = -25, y = 81, datetime = 1900-01-07 09:00:00, uid = cQFUsvMk)\n",
+ " * SpaceTimeRecord(x = -116, y = 43, datetime = 1900-01-09 01:00:00, uid = MDpcWsK8)\n",
+ " - with children:\n",
+ " OctTree:\n",
+ " - boundary: SpaceTimeRectangle(west=-180, east=-90.0, south=45.0, north=90, start=datetime.datetime(1900, 1, 1, 0, 0), end=datetime.datetime(1900, 1, 8, 17, 45))\n",
+ " - capacity: 10\n",
+ " - depth: 2\n",
+ " - max_depth: 25\n",
+ " - contents:\n",
+ " - number of elements: 10\n",
+ " * SpaceTimeRecord(x = -130, y = 80, datetime = 1900-01-01 03:00:00, uid = sd0nBvvS)\n",
+ " * SpaceTimeRecord(x = -148, y = 78, datetime = 1900-01-06 03:00:00, uid = FgJRfXD9)\n",
+ " * SpaceTimeRecord(x = -153, y = 58, datetime = 1900-01-03 12:00:00, uid = AHWomxBm)\n",
+ " * SpaceTimeRecord(x = -160, y = 47, datetime = 1900-01-06 18:00:00, uid = 3p50Ejkq)\n",
+ " * SpaceTimeRecord(x = -91, y = 60, datetime = 1900-01-07 06:00:00, uid = 1Psbg1Vk)\n",
+ " * SpaceTimeRecord(x = -138, y = 54, datetime = 1900-01-08 10:00:00, uid = kDwksPIp)\n",
+ " * SpaceTimeRecord(x = -99, y = 86, datetime = 1900-01-05 12:00:00, uid = gfhX01rL)\n",
+ " * SpaceTimeRecord(x = -96, y = 54, datetime = 1900-01-04 23:00:00, uid = o7lz8pja)\n",
+ " * SpaceTimeRecord(x = -163, y = 79, datetime = 1900-01-07 22:00:00, uid = 2Fw915S3)\n",
+ " * SpaceTimeRecord(x = -155, y = 74, datetime = 1900-01-08 09:00:00, uid = 9pL97BD0)\n",
+ " - with children:\n",
+ " OctTree:\n",
+ " - boundary: SpaceTimeRectangle(west=-180, east=-135.0, south=67.5, north=90, start=datetime.datetime(1900, 1, 1, 0, 0), end=datetime.datetime(1900, 1, 4, 20, 52, 30))\n",
+ " - capacity: 10\n",
+ " - depth: 3\n",
+ " - max_depth: 25\n",
+ " - contents:\n",
+ " - number of elements: 10\n",
+ " * SpaceTimeRecord(x = -173, y = 71, datetime = 1900-01-04 03:00:00, uid = ThLEI8lF)\n",
+ " * SpaceTimeRecord(x = -167, y = 83, datetime = 1900-01-04 03:00:00, uid = Q5FzwxD5)\n",
+ " * SpaceTimeRecord(x = -167, y = 88, datetime = 1900-01-01 16:00:00, uid = DoCBI1YI)\n",
+ " * SpaceTimeRecord(x = -141, y = 80, datetime = 1900-01-03 16:00:00, uid = 01SVlWsE)\n",
+ " * SpaceTimeRecord(x = -135, y = 68, datetime = 1900-01-03 22:00:00, uid = Jx2uI4Op)\n",
+ " * SpaceTimeRecord(x = -163, y = 77, datetime = 1900-01-03 21:00:00, uid = DoOKHLix)\n",
+ " * SpaceTimeRecord(x = -157, y = 84, datetime = 1900-01-02 11:00:00, uid = lXiFUOBn)\n",
+ " * SpaceTimeRecord(x = -145, y = 78, datetime = 1900-01-02 05:00:00, uid = 3ngKJmcS)\n",
+ " * SpaceTimeRecord(x = -179, y = 89, datetime = 1900-01-04 01:00:00, uid = KQXXjSTT)\n",
+ " * SpaceTimeRecord(x = -171, y = 80, datetime = 1900-01-04 19:00:00, uid = znugCZWi)\n",
+ " - with children:\n",
+ " OctTree:\n",
+ " - boundary: SpaceTimeRectangle(west=-180, east=-157.5, south=78.75, north=90, start=datetime.datetime(1900, 1, 1, 0, 0), end=datetime.datetime(1900, 1, 2, 22, 26, 15))\n",
+ " - capacity: 10\n",
+ " - depth: 4\n",
+ " - max_depth: 25\n",
+ " - contents:\n",
+ " - number of elements: 10\n",
+ " * SpaceTimeRecord(x = -164, y = 79, datetime = 1900-01-01 18:00:00, uid = L0scr6Dw)\n",
+ " * SpaceTimeRecord(x = -165, y = 87, datetime = 1900-01-02 12:00:00, uid = P2JSVMig)\n",
+ " * SpaceTimeRecord(x = -158, y = 80, datetime = 1900-01-01 07:00:00, uid = rrfLnl9a000)\n",
+ " * SpaceTimeRecord(x = -179, y = 88, datetime = 1900-01-01 18:00:00, uid = piKfH7lZ000)\n",
+ " * SpaceTimeRecord(x = -162, y = 85, datetime = 1900-01-01 00:00:00, uid = TzzMqFl4000)\n",
+ " * SpaceTimeRecord(x = -170, y = 85, datetime = 1900-01-01 02:00:00, uid = v6BVfkmP000)\n",
+ " * SpaceTimeRecord(x = -177, y = 87, datetime = 1900-01-01 08:00:00, uid = sAxKIWXJ001)\n",
+ " * SpaceTimeRecord(x = -164, y = 79, datetime = 1900-01-01 02:00:00, uid = 55vXE5De001)\n",
+ " * SpaceTimeRecord(x = -167, y = 87, datetime = 1900-01-02 00:00:00, uid = RzUJ5Q7h001)\n",
+ " * SpaceTimeRecord(x = -170, y = 89, datetime = 1900-01-01 20:00:00, uid = ipKiaytp002)\n",
+ " - with children:\n",
+ " OctTree:\n",
+ " - boundary: SpaceTimeRectangle(west=-180, east=-168.75, south=84.375, north=90, start=datetime.datetime(1900, 1, 1, 0, 0), end=datetime.datetime(1900, 1, 1, 23, 13, 7, 500000))\n",
+ " - capacity: 10\n",
+ " - depth: 5\n",
+ " - max_depth: 25\n",
+ " - contents:\n",
+ " - number of elements: 10\n",
+ " * SpaceTimeRecord(x = -178, y = 87, datetime = 1900-01-01 11:00:00, uid = 9i3tUAKH003)\n",
+ " * SpaceTimeRecord(x = -169, y = 88, datetime = 1900-01-01 23:00:00, uid = ib1mXyZJ003)\n",
+ " * SpaceTimeRecord(x = -174, y = 88, datetime = 1900-01-01 03:00:00, uid = vYJ8DamM004)\n"
]
}
],
"source": [
- "%%timeit\n",
- "rec = random.choice(test_recs)\n",
- "nearby_ships(lon=rec.lon, lat=rec.lat, dt=rec.datetime, max_dist=dist, dt_gap=dt, date_col=\"datetime\", pool=df, filter_datetime=True)"
+ "s = str(otree)\n",
+ "print(\"\\n\".join(s.split(\"\\n\")[:100]))"
]
},
{
"cell_type": "markdown",
- "id": "d148f129-9d8c-4c46-8f01-3e9c1e93e81a",
+ "id": "6b02c2ea-6566-47c2-97e0-43d8b18e0713",
"metadata": {},
"source": [
- "## Verify\n",
+ "## Time Execution\n",
"\n",
- "Check that records are the same"
+ "Testing the identification of nearby points against the original full search"
]
},
{
"cell_type": "code",
- "execution_count": 12,
- "id": "11f3d73a-fbe5-4f27-88d8-d0d687bd0eac",
+ "execution_count": 9,
+ "id": "094b588c-e938-4838-9719-1defdfff74fa",
"metadata": {},
+ "outputs": [],
+ "source": [
+ "dts = pl.datetime_range(\n",
+ " datetime(1900, 1, 1),\n",
+ " datetime(1900, 2, 1),\n",
+ " interval=\"1h\",\n",
+ " eager=True,\n",
+ " closed=\"left\",\n",
+ ")\n",
+ "N = dts.len()\n",
+ "lons = 180 - 360 * np.random.rand(N)\n",
+ "lats = 90 - 180 * np.random.rand(N)\n",
+ "test_df = pl.DataFrame({\"lon\": lons, \"lat\": lats, \"datetime\": dts})\n",
+ "test_recs = [Record(*r) for r in test_df.rows()]\n",
+ "dt = timedelta(days=1)\n",
+ "dist = 350"
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": 10,
+ "id": "66a48b86-d449-45d2-9837-2b3e07f5563d",
+ "metadata": {},
+ "outputs": [
+ {
+ "name": "stdout",
+ "output_type": "stream",
+ "text": [
+ "1.32 ms ± 20.8 μs per loop (mean ± std. dev. of 7 runs, 1,000 loops each)\n"
+ ]
+ }
+ ],
+ "source": [
+ "%%timeit\n",
+ "otree.nearby_points(random.choice(test_recs), dist=dist, t_dist=dt)"
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": 11,
+ "id": "8b9279ed-6f89-4423-8833-acd0b365eb7b",
+ "metadata": {
+ "scrolled": true
+ },
"outputs": [
{
"name": "stdout",
"output_type": "stream",
"text": [
- "CPU times: user 2.52 s, sys: 253 ms, total: 2.78 s\n",
- "Wall time: 2.66 s\n"
+ "12.1 ms ± 81.4 μs per loop (mean ± std. dev. of 7 runs, 100 loops each)\n"
]
}
],
+ "source": [
+ "%%timeit\n",
+ "rec = random.choice(test_recs)\n",
+ "nearby_ships(\n",
+ " lon=rec.lon,\n",
+ " lat=rec.lat,\n",
+ " dt=rec.datetime,\n",
+ " max_dist=dist,\n",
+ " dt_gap=dt,\n",
+ " date_col=\"datetime\",\n",
+ " pool=df,\n",
+ " filter_datetime=True,\n",
+ ")"
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "id": "d148f129-9d8c-4c46-8f01-3e9c1e93e81a",
+ "metadata": {},
+ "source": [
+ "## Verify\n",
+ "\n",
+ "Check that records are the same"
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": null,
+ "id": "11f3d73a-fbe5-4f27-88d8-d0d687bd0eac",
+ "metadata": {},
+ "outputs": [],
"source": [
"%%time\n",
"dist = 250\n",
"for _ in range(250):\n",
" rec = Record(*random.choice(df.rows()))\n",
- " orig = nearby_ships(lon=rec.lon, lat=rec.lat, dt=rec.datetime, max_dist=dist, dt_gap=dt, date_col=\"datetime\", pool=df, filter_datetime=True)\n",
+ " orig = nearby_ships(lon=rec.lon, lat=rec.lat, dt=rec.datetime, max_dist=dist, dt_gap=dt, date_col=\"datetime\", pool=df, filter_datetime=True) # noqa\n",
" tree = otree.nearby_points(rec, dist=dist, t_dist=dt)\n",
" if orig.height > 0:\n",
" if not tree:\n",
@@ -649,7 +662,7 @@
" print(\"NO TREE!\")\n",
" print(f\"{orig = }\")\n",
" else:\n",
- " tree = pl.from_records([(r.lon, r.lat, r.datetime, r.uid) for r in tree], orient=\"row\").rename({\"column_0\": \"lon\", \"column_1\": \"lat\", \"column_2\": \"datetime\", \"column_3\": \"uid\"})\n",
+ " tree = pl.from_records([(r.lon, r.lat, r.datetime, r.uid) for r in tree], orient=\"row\").rename({\"column_0\": \"lon\", \"column_1\": \"lat\", \"column_2\": \"datetime\", \"column_3\": \"uid\"}) # noqa\n",
" if tree.height != orig.height:\n",
" print(\"Tree and Orig Heights Do Not Match\")\n",
" print(f\"{orig = }\")\n",
@@ -672,12 +685,16 @@
},
{
"cell_type": "code",
- "execution_count": 13,
+ "execution_count": null,
"id": "4c392292-2d9f-4301-afb5-019fde069a1e",
"metadata": {},
"outputs": [],
"source": [
- "out = otree.nearby_points(Record(179.5, -43.1, datetime(1900, 1, 14, 13)), dist=200, t_dist=timedelta(days=3))\n",
+ "out = otree.nearby_points(\n",
+ " Record(179.5, -43.1, datetime(1900, 1, 14, 13)),\n",
+ " dist=200,\n",
+ " t_dist=timedelta(days=3),\n",
+ ")\n",
"for o in out:\n",
" print(o)"
]
@@ -685,7 +702,7 @@
],
"metadata": {
"kernelspec": {
- "display_name": "GeoSpatialTools",
+ "display_name": "geospatialtools",
"language": "python",
"name": "geospatialtools"
},
@@ -699,7 +716,7 @@
"name": "python",
"nbconvert_exporter": "python",
"pygments_lexer": "ipython3",
- "version": "3.12.7"
+ "version": "3.11.11"
}
},
"nbformat": 4,