Skip to content

Commit

Permalink
symmetric kernel experimentation notebook (#9)
Browse files Browse the repository at this point in the history
  • Loading branch information
@keewis
keewis authored May 28, 2024
1 parent ccf49ad commit 348badf
Showing 1 changed file with 394 additions and 0 deletions.
394 changes: 394 additions & 0 deletions experimentation/kernel.ipynb
View file
Original file line number Diff line number Diff line change
@@ -0,0 +1,394 @@
{
"cells": [
{
"cell_type": "code",
"execution_count": null,
"id": "0",
"metadata": {},
"outputs": [],
"source": [
"import dask.array as da\n",
"import healpy as hp\n",
"import numpy as np\n",
"\n",
"import healpix_convolution as hc"
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "1",
"metadata": {},
"outputs": [],
"source": [
"import folium\n",
"import matplotlib.pyplot as plt"
]
},
{
"cell_type": "markdown",
"id": "2",
"metadata": {
"jp-MarkdownHeadingCollapsed": true
},
"source": [
"## experimentation"
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "3",
"metadata": {},
"outputs": [],
"source": [
"resolution = 4\n",
"cell_ids = np.arange(12 * 4**resolution)\n",
"indexing_scheme = \"nested\""
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "4",
"metadata": {},
"outputs": [],
"source": [
"truncate = 4\n",
"sigma = 0.1 # in radians\n",
"cell_distance = hp.nside2resol(2**resolution, arcmin=False)\n",
"ring = int((truncate * sigma / cell_distance) // 2)\n",
"\n",
"neighbours = hc.neighbours(\n",
" cell_ids, resolution=resolution, indexing_scheme=indexing_scheme, ring=ring\n",
")\n",
"distances = hc.angular_distances(\n",
" neighbours, resolution=resolution, indexing_scheme=indexing_scheme\n",
")\n",
"mask = neighbours == -1\n",
"\n",
"sigma2 = sigma * sigma\n",
"phi_x = np.where(mask, 0, np.exp(-0.5 / sigma2 * distances**2))\n",
"kernel = phi_x / phi_x.sum(axis=-1)[:, None]\n",
"kernel.shape"
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "5",
"metadata": {},
"outputs": [],
"source": [
"import sparse"
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "6",
"metadata": {},
"outputs": [],
"source": [
"mask = np.reshape(neighbours, -1) != -1\n",
"coords = np.reshape(\n",
" np.stack(\n",
" [\n",
" np.repeat(cell_ids[:, None], repeats=neighbours.shape[-1], axis=-1),\n",
" neighbours,\n",
" ],\n",
" axis=0,\n",
" ),\n",
" (2, -1),\n",
")\n",
"\n",
"kernel_ = np.reshape(kernel, -1)[mask]\n",
"coords_ = np.reshape(coords, (2, -1))[:, mask]"
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "7",
"metadata": {},
"outputs": [],
"source": [
"kernel_matrix = sparse.COO(\n",
" data=kernel_, coords=coords_, shape=(cell_ids.size, cell_ids.size), fill_value=0\n",
")\n",
"kernel_matrix"
]
},
{
"cell_type": "markdown",
"id": "8",
"metadata": {},
"source": [
"## dask awareness"
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "9",
"metadata": {},
"outputs": [],
"source": [
"resolution = 4\n",
"kernel_size = 3\n",
"indexing_scheme = \"ring\"\n",
"sigma = 0.1\n",
"\n",
"cell_ids = da.arange(12 * 4**resolution, chunks=(1000,))\n",
"cell_ids"
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "10",
"metadata": {},
"outputs": [],
"source": [
"cell_ids_ = np.reshape(cell_ids, (-1,))\n",
"\n",
"# TODO: figure out whether there is a better way of defining the units of `sigma`\n",
"if kernel_size is not None:\n",
" ring = int(kernel_size / 2)\n",
"else:\n",
" cell_distance = hp.nside2resol(2**resolution, arcmin=False)\n",
" ring = int((truncate * sigma / cell_distance) // 2)\n",
"\n",
"nb = hc.neighbours(\n",
" cell_ids_, resolution=resolution, indexing_scheme=indexing_scheme, ring=ring\n",
")\n",
"d = hc.angular_distances(nb, resolution=resolution, indexing_scheme=indexing_scheme)\n",
"\n",
"sigma2 = sigma * sigma\n",
"phi_x = np.exp(-0.5 / sigma2 * d**2)\n",
"masked = np.where(nb == -1, 0, phi_x)\n",
"normalized = masked / np.sum(masked, axis=1, keepdims=True)\n",
"normalized"
]
},
{
"cell_type": "raw",
"id": "11",
"metadata": {},
"source": [
"da.map_blocks?"
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "12",
"metadata": {},
"outputs": [],
"source": [
"import sparse"
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "13",
"metadata": {},
"outputs": [],
"source": [
"cell_ids__ = np.repeat(cell_ids_[:, None], axis=-1, repeats=nb.shape[1])\n",
"cell_ids__"
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "14",
"metadata": {},
"outputs": [],
"source": [
"?da.map_blocks"
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "15",
"metadata": {},
"outputs": [],
"source": [
"cell_ids__.chunks"
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "16",
"metadata": {},
"outputs": [],
"source": [
"shape = (1000, cell_ids.size)\n",
"matrix = da.map_blocks(\n",
" hc.kernels.common.create_sparse,\n",
" cell_ids__,\n",
" nb,\n",
" normalized,\n",
" shape=shape,\n",
" meta=sparse.COO.from_numpy(np.array((), dtype=\"float64\")),\n",
" drop_axis=1,\n",
" new_axis=1,\n",
" chunks=(cell_ids__.chunks[0], cell_ids.size),\n",
")\n",
"matrix"
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "17",
"metadata": {},
"outputs": [],
"source": [
"display(cell_ids__, nb, normalized)"
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "18",
"metadata": {},
"outputs": [],
"source": [
"matrix.compute()"
]
},
{
"cell_type": "markdown",
"id": "19",
"metadata": {},
"source": [
"## module version"
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "20",
"metadata": {},
"outputs": [],
"source": [
"resolution = 3\n",
"cell_ids = np.arange(12 * 4**resolution)\n",
"indexing_scheme = \"nested\"\n",
"sigma = 0.1\n",
"truncate = 4.0"
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "21",
"metadata": {},
"outputs": [],
"source": [
"%%time\n",
"kernel = hc.kernels.gaussian_kernel(\n",
" cell_ids,\n",
" resolution=resolution,\n",
" indexing_scheme=indexing_scheme,\n",
" sigma=sigma,\n",
" truncate=truncate,\n",
")\n",
"kernel"
]
},
{
"cell_type": "markdown",
"id": "22",
"metadata": {},
"source": [
"### verification"
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "23",
"metadata": {},
"outputs": [],
"source": [
"norm = np.sum(kernel, axis=1).todense()\n",
"norm"
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "24",
"metadata": {},
"outputs": [],
"source": [
"fig, ax = plt.subplots(figsize=(14, 14))\n",
"\n",
"mappable = ax.imshow(kernel.todense())\n",
"fig.colorbar(mappable)"
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "25",
"metadata": {},
"outputs": [],
"source": [
"kernel_ = kernel[0, :].todense()"
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "26",
"metadata": {},
"outputs": [],
"source": [
"import healpy as hp"
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "27",
"metadata": {},
"outputs": [],
"source": [
"hp.newvisufunc.projview(kernel_, nest=True)"
]
},
{
"cell_type": "markdown",
"id": "28",
"metadata": {},
"source": [
"- subdomain convolution\n",
"- image pyramid (up/downgrading)\n",
"- neighbour ordering\n",
"- chunked kernel"
]
}
],
"metadata": {
"language_info": {
"codemirror_mode": {
"name": "ipython",
"version": 3
},
"file_extension": ".py",
"mimetype": "text/x-python",
"name": "python",
"nbconvert_exporter": "python",
"pygments_lexer": "ipython3"
}
},
"nbformat": 4,
"nbformat_minor": 5
}

0 comments on commit 348badf

Please sign in to comment.