grid.html 128.16 KiB
<!DOCTYPE html>
<html lang="en" data-content_root="../../" >
<head>
<meta charset="utf-8" />
<meta name="viewport" content="width=device-width, initial-scale=1.0" />
<title>arbdmodel.grid — ARBD Model Documentation</title>
<script data-cfasync="false">
document.documentElement.dataset.mode = localStorage.getItem("mode") || "";
document.documentElement.dataset.theme = localStorage.getItem("theme") || "";
</script>
<!-- Loaded before other Sphinx assets -->
<link href="../../_static/styles/theme.css?digest=dfe6caa3a7d634c4db9b" rel="stylesheet" />
<link href="../../_static/styles/bootstrap.css?digest=dfe6caa3a7d634c4db9b" rel="stylesheet" />
<link href="../../_static/styles/pydata-sphinx-theme.css?digest=dfe6caa3a7d634c4db9b" rel="stylesheet" />
<link href="../../_static/vendor/fontawesome/6.5.2/css/all.min.css?digest=dfe6caa3a7d634c4db9b" rel="stylesheet" />
<link rel="preload" as="font" type="font/woff2" crossorigin href="../../_static/vendor/fontawesome/6.5.2/webfonts/fa-solid-900.woff2" />
<link rel="preload" as="font" type="font/woff2" crossorigin href="../../_static/vendor/fontawesome/6.5.2/webfonts/fa-brands-400.woff2" />
<link rel="preload" as="font" type="font/woff2" crossorigin href="../../_static/vendor/fontawesome/6.5.2/webfonts/fa-regular-400.woff2" />
<link rel="stylesheet" type="text/css" href="../../_static/pygments.css?v=03e43079" />
<link rel="stylesheet" type="text/css" href="../../_static/styles/sphinx-book-theme.css?v=eba8b062" />
<link rel="stylesheet" type="text/css" href="../../_static/togglebutton.css?v=13237357" />
<link rel="stylesheet" type="text/css" href="../../_static/copybutton.css?v=76b2166b" />
<link rel="stylesheet" type="text/css" href="../../_static/mystnb.4510f1fc1dee50b3e5859aac5469c37c29e427902b24a333a5f9fcb2f0b3ac41.css?v=be8a1c11" />
<link rel="stylesheet" type="text/css" href="../../_static/sphinx-thebe.css?v=4fa983c6" />
<link rel="stylesheet" type="text/css" href="../../_static/tabs.css?v=4c969af8" />
<link rel="stylesheet" type="text/css" href="../../_static/proof.css?v=b4b7a797" />
<link rel="stylesheet" type="text/css" href="../../_static/styles/sphinx-examples.css?v=e236af4b" />
<link rel="stylesheet" type="text/css" href="../../_static/css/tooltipster.custom.css?v=7bc2f056" />
<link rel="stylesheet" type="text/css" href="../../_static/css/tooltipster.bundle.min.css?v=37217874" />
<link rel="stylesheet" type="text/css" href="../../_static/css/tooltipster-sideTip-shadow.min.css?v=6227e517" />
<link rel="stylesheet" type="text/css" href="../../_static/css/tooltipster-sideTip-punk.min.css?v=94669e23" />
<link rel="stylesheet" type="text/css" href="../../_static/css/tooltipster-sideTip-noir.min.css?v=21a39f42" />
<link rel="stylesheet" type="text/css" href="../../_static/css/tooltipster-sideTip-light.min.css?v=a18b2449" />
<link rel="stylesheet" type="text/css" href="../../_static/css/tooltipster-sideTip-borderless.min.css?v=dbff53e4" />
<link rel="stylesheet" type="text/css" href="../../_static/css/micromodal.css?v=d7bf34ee" />
<link rel="stylesheet" type="text/css" href="../../_static/sphinx-design.min.css?v=95c83b7e" />
<!-- Pre-loaded scripts that we'll load fully later -->
<link rel="preload" as="script" href="../../_static/scripts/bootstrap.js?digest=dfe6caa3a7d634c4db9b" />
<link rel="preload" as="script" href="../../_static/scripts/pydata-sphinx-theme.js?digest=dfe6caa3a7d634c4db9b" />
<script src="../../_static/vendor/fontawesome/6.5.2/js/all.min.js?digest=dfe6caa3a7d634c4db9b"></script>
<script src="../../_static/jquery.js?v=5d32c60e"></script>
<script src="../../_static/_sphinx_javascript_frameworks_compat.js?v=2cd50e6c"></script>
<script src="../../_static/documentation_options.js?v=9eb32ce0"></script>
<script src="../../_static/doctools.js?v=9a2dae69"></script>
<script src="../../_static/sphinx_highlight.js?v=dc90522c"></script>
<script src="../../_static/clipboard.min.js?v=a7894cd8"></script>
<script src="../../_static/copybutton.js?v=f281be69"></script>
<script src="../../_static/scripts/sphinx-book-theme.js?v=887ef09a"></script>
<script src="../../_static/tabs.js?v=3ee01567"></script>
<script src="../../_static/js/hoverxref.js?v=c95ade4f"></script>
<script src="../../_static/js/tooltipster.bundle.min.js?v=18bf091b"></script>
<script src="../../_static/js/micromodal.min.js?v=04d6302d"></script>
<script>let toggleHintShow = 'Click to show';</script>
<script>let toggleHintHide = 'Click to hide';</script>
<script>let toggleOpenOnPrint = 'true';</script>
<script src="../../_static/togglebutton.js?v=4a39c7ea"></script>
<script>var togglebuttonSelector = '.toggle, .admonition.dropdown';</script> <script src="../../_static/design-tabs.js?v=f930bc37"></script>
<script>const THEBE_JS_URL = "https://unpkg.com/thebe@0.8.2/lib/index.js"; const thebe_selector = ".thebe,.cell"; const thebe_selector_input = "pre"; const thebe_selector_output = ".output, .cell_output"</script>
<script async="async" src="../../_static/sphinx-thebe.js?v=c100c467"></script>
<script>var togglebuttonSelector = '.toggle, .admonition.dropdown';</script>
<script>const THEBE_JS_URL = "https://unpkg.com/thebe@0.8.2/lib/index.js"; const thebe_selector = ".thebe,.cell"; const thebe_selector_input = "pre"; const thebe_selector_output = ".output, .cell_output"</script>
<script>window.MathJax = {"options": {"processHtmlClass": "tex2jax_process|mathjax_process|math|output_area"}}</script>
<script defer="defer" src="https://cdn.jsdelivr.net/npm/mathjax@3/es5/tex-mml-chtml.js"></script>
<script>DOCUMENTATION_OPTIONS.pagename = '_modules/arbdmodel/grid';</script>
<link rel="index" title="Index" href="../../genindex.html" />
<link rel="search" title="Search" href="../../search.html" />
<meta name="viewport" content="width=device-width, initial-scale=1"/>
<meta name="docsearch:language" content="en"/>
</head>
<body data-bs-spy="scroll" data-bs-target=".bd-toc-nav" data-offset="180" data-bs-root-margin="0px 0px -60%" data-default-mode="">
<div id="pst-skip-link" class="skip-link d-print-none"><a href="#main-content">Skip to main content</a></div>
<div id="pst-scroll-pixel-helper"></div>
<button type="button" class="btn rounded-pill" id="pst-back-to-top">
<i class="fa-solid fa-arrow-up"></i>Back to top</button>
<input type="checkbox"
class="sidebar-toggle"
id="pst-primary-sidebar-checkbox"/>
<label class="overlay overlay-primary" for="pst-primary-sidebar-checkbox"></label>
<input type="checkbox"
class="sidebar-toggle"
id="pst-secondary-sidebar-checkbox"/>
<label class="overlay overlay-secondary" for="pst-secondary-sidebar-checkbox"></label>
<div class="search-button__wrapper">
<div class="search-button__overlay"></div>
<div class="search-button__search-container">
<form class="bd-search d-flex align-items-center"
action="../../search.html"
method="get">
<i class="fa-solid fa-magnifying-glass"></i>
<input type="search"
class="form-control"
name="q"
id="search-input"
placeholder="Search..."
aria-label="Search..."
autocomplete="off"
autocorrect="off"
autocapitalize="off"
spellcheck="false"/>
<span class="search-button__kbd-shortcut"><kbd class="kbd-shortcut__modifier">Ctrl</kbd>+<kbd>K</kbd></span>
</form></div>
</div>
<div class="pst-async-banner-revealer d-none">
<aside id="bd-header-version-warning" class="d-none d-print-none" aria-label="Version warning"></aside>
</div>
<header class="bd-header navbar navbar-expand-lg bd-navbar d-print-none">
</header>
<div class="bd-container">
<div class="bd-container__inner bd-page-width">
<div class="bd-sidebar-primary bd-sidebar">
<div class="sidebar-header-items sidebar-primary__section">
</div>
<div class="sidebar-primary-items__start sidebar-primary__section">
<div class="sidebar-primary-item">
<a class="navbar-brand logo" href="../../intro.html">
<p class="title logo__title">ARBD Model Documentation</p>
</a></div>
<div class="sidebar-primary-item">
<script>
document.write(`
<button class="btn search-button-field search-button__button" title="Search" aria-label="Search" data-bs-placement="bottom" data-bs-toggle="tooltip">
<i class="fa-solid fa-magnifying-glass"></i>
<span class="search-button__default-text">Search</span>
<span class="search-button__kbd-shortcut"><kbd class="kbd-shortcut__modifier">Ctrl</kbd>+<kbd class="kbd-shortcut__modifier">K</kbd></span>
</button>
`);
</script></div>
<div class="sidebar-primary-item"><nav class="bd-links bd-docs-nav" aria-label="Main">
<div class="bd-toc-item navbar-nav active">
<p aria-level="2" class="caption" role="heading"><span class="caption-text">Getting Started</span></p>
<ul class="nav bd-sidenav">
<li class="toctree-l1"><a class="reference internal" href="../../tutorials/index.html">Tutorials</a></li>
</ul>
<p aria-level="2" class="caption" role="heading"><span class="caption-text">API Reference</span></p>
<ul class="nav bd-sidenav">
<li class="toctree-l1"><a class="reference internal" href="../../api/index.html">API Reference</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../api/core/index.html">Core</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../api/interaction_potentials/index.html">Interaction Potentials</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../api/polymer_modeling/index.html">Polymer Modeling</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../api/rigidbody_models/index.html">RigidBody Models</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../api/shape-based_models/index.html">Shape-Based Models</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../api/simulation_engines/index.html">Simulation Engines</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../api/utilities/index.html">Utilities</a></li>
</ul>
</div>
</nav></div>
</div>
<div class="sidebar-primary-items__end sidebar-primary__section">
</div>
<div id="rtd-footer-container"></div>
</div>
<main id="main-content" class="bd-main" role="main">
<div class="sbt-scroll-pixel-helper"></div>
<div class="bd-content">
<div class="bd-article-container">
<div class="bd-header-article d-print-none">
<div class="header-article-items header-article__inner">
<div class="header-article-items__start">
<div class="header-article-item"><button class="sidebar-toggle primary-toggle btn btn-sm" title="Toggle primary sidebar" data-bs-placement="bottom" data-bs-toggle="tooltip">
<span class="fa-solid fa-bars"></span>
</button></div>
</div>
<div class="header-article-items__end">
<div class="header-article-item">
<div class="article-header-buttons">
<button onclick="toggleFullScreen()"
class="btn btn-sm btn-fullscreen-button"
title="Fullscreen mode"
data-bs-placement="bottom" data-bs-toggle="tooltip"
>
<span class="btn__icon-container">
<i class="fas fa-expand"></i>
</span>
</button>
<script>
document.write(`
<button class="btn btn-sm nav-link pst-navbar-icon theme-switch-button" title="light/dark" aria-label="light/dark" data-bs-placement="bottom" data-bs-toggle="tooltip">
<i class="theme-switch fa-solid fa-sun fa-lg" data-mode="light"></i>
<i class="theme-switch fa-solid fa-moon fa-lg" data-mode="dark"></i>
<i class="theme-switch fa-solid fa-circle-half-stroke fa-lg" data-mode="auto"></i>
</button>
`);
</script>
<script>
document.write(`
<button class="btn btn-sm pst-navbar-icon search-button search-button__button" title="Search" aria-label="Search" data-bs-placement="bottom" data-bs-toggle="tooltip">
<i class="fa-solid fa-magnifying-glass fa-lg"></i>
</button>
`);
</script>
</div></div>
</div>
</div>
</div>
<div id="jb-print-docs-body" class="onlyprint">
<h1></h1>
<!-- Table of contents -->
<div id="print-main-content">
<div id="jb-print-toc">
</div>
</div>
</div>
<div id="searchbox"></div>
<article class="bd-article">
<h1>Source code for arbdmodel.grid</h1><div class="highlight"><pre>
<span></span><span class="kn">from</span><span class="w"> </span><span class="nn">__future__</span><span class="w"> </span><span class="kn">import</span> <span class="n">absolute_import</span><span class="p">,</span> <span class="n">print_function</span>
<span class="kn">import</span><span class="w"> </span><span class="nn">numpy</span><span class="w"> </span><span class="k">as</span><span class="w"> </span><span class="nn">np</span>
<span class="kn">from</span><span class="w"> </span><span class="nn">scipy</span><span class="w"> </span><span class="kn">import</span> <span class="n">signal</span>
<span class="kn">import</span><span class="w"> </span><span class="nn">os</span><span class="o">,</span><span class="nn">sys</span>
<span class="kn">from</span><span class="w"> </span><span class="nn">.logger</span><span class="w"> </span><span class="kn">import</span> <span class="n">logger</span>
<span class="kn">import</span><span class="w"> </span><span class="nn">unittest</span>
<span class="kn">import</span><span class="w"> </span><span class="nn">numpy</span><span class="w"> </span><span class="k">as</span><span class="w"> </span><span class="nn">np</span>
<span class="kn">from</span><span class="w"> </span><span class="nn">pathlib</span><span class="w"> </span><span class="kn">import</span> <span class="n">Path</span>
<div class="viewcode-block" id="writeDx">
<a class="viewcode-back" href="../../api/utilities/grid.html#arbdmodel.grid.writeDx">[docs]</a>
<span class="k">def</span><span class="w"> </span><span class="nf">writeDx</span><span class="p">(</span><span class="n">outfile</span><span class="p">,</span> <span class="n">data</span><span class="p">,</span> <span class="n">origin</span><span class="p">,</span> <span class="n">delta</span><span class="p">,</span> <span class="n">fmt</span><span class="o">=</span><span class="s2">"</span><span class="si">%.12f</span><span class="s2">"</span><span class="p">):</span>
<span class="w"> </span><span class="sd">"""Write 3D grid data to an OpenDX format file.</span>
<span class="sd"> Write 3D grid data to an OpenDX format file.</span>
<span class="sd"> </span>
<span class="sd"> OpenDX is a visualization software format that can be used to visualize </span>
<span class="sd"> volumetric data.</span>
<span class="sd"> </span>
<span class="sd"> Parameters</span>
<span class="sd"> ----------</span>
<span class="sd"> outfile : str or file-like object</span>
<span class="sd"> The path to the output file or a file-like object.</span>
<span class="sd"> data : numpy.ndarray</span>
<span class="sd"> 3D array containing the grid data to be written.</span>
<span class="sd"> origin : list or tuple</span>
<span class="sd"> The (x, y, z) coordinates of the origin of the grid.</span>
<span class="sd"> delta : list or tuple</span>
<span class="sd"> The grid spacing in the (x, y, z) directions.</span>
<span class="sd"> fmt : str, optional</span>
<span class="sd"> Format string for the data values. Default is "%.12f".</span>
<span class="sd"> </span>
<span class="sd"> Note</span>
<span class="sd"> -----</span>
<span class="sd"> The output file follows the OpenDX format specification:</span>
<span class="sd"> http://opendx.sdsc.edu/docs/html/pages/usrgu068.htm#HDREDF</span>
<span class="sd"> </span>
<span class="sd"> The data is written in a way that is compatible with visualization software</span>
<span class="sd"> that accepts OpenDX format.</span>
<span class="sd"> </span>
<span class="sd"> Raises</span>
<span class="sd"> ------</span>
<span class="sd"> AssertionError</span>
<span class="sd"> If data is not a 3D array or if origin or delta do not have length 3.</span>
<span class="sd"> """</span>
<span class="n">shape</span> <span class="o">=</span> <span class="n">np</span><span class="o">.</span><span class="n">shape</span><span class="p">(</span><span class="n">data</span><span class="p">)</span>
<span class="n">num</span> <span class="o">=</span> <span class="n">np</span><span class="o">.</span><span class="n">prod</span><span class="p">(</span><span class="n">shape</span><span class="p">)</span> <span class="k">assert</span><span class="p">(</span> <span class="nb">len</span><span class="p">(</span><span class="n">shape</span><span class="p">)</span> <span class="o">==</span> <span class="mi">3</span> <span class="p">)</span>
<span class="k">assert</span><span class="p">(</span> <span class="nb">len</span><span class="p">(</span><span class="n">origin</span><span class="p">)</span> <span class="o">==</span> <span class="mi">3</span> <span class="p">)</span>
<span class="k">assert</span><span class="p">(</span> <span class="nb">len</span><span class="p">(</span><span class="n">delta</span><span class="p">)</span> <span class="o">==</span> <span class="mi">3</span> <span class="p">)</span>
<span class="n">headerInfo</span> <span class="o">=</span> <span class="nb">dict</span><span class="p">(</span> <span class="n">nx</span><span class="o">=</span><span class="n">shape</span><span class="p">[</span><span class="mi">0</span><span class="p">],</span> <span class="n">ny</span><span class="o">=</span><span class="n">shape</span><span class="p">[</span><span class="mi">1</span><span class="p">],</span> <span class="n">nz</span><span class="o">=</span><span class="n">shape</span><span class="p">[</span><span class="mi">2</span><span class="p">],</span>
<span class="n">ox</span><span class="o">=</span><span class="n">origin</span><span class="p">[</span><span class="mi">0</span><span class="p">],</span> <span class="n">oy</span><span class="o">=</span><span class="n">origin</span><span class="p">[</span><span class="mi">1</span><span class="p">],</span> <span class="n">oz</span><span class="o">=</span><span class="n">origin</span><span class="p">[</span><span class="mi">2</span><span class="p">],</span>
<span class="n">dx</span><span class="o">=</span><span class="n">delta</span><span class="p">[</span><span class="mi">0</span><span class="p">],</span> <span class="n">dy</span><span class="o">=</span><span class="n">delta</span><span class="p">[</span><span class="mi">1</span><span class="p">],</span> <span class="n">dz</span><span class="o">=</span><span class="n">delta</span><span class="p">[</span><span class="mi">2</span><span class="p">],</span>
<span class="n">num</span><span class="o">=</span><span class="n">num</span>
<span class="p">)</span>
<span class="n">data</span> <span class="o">=</span> <span class="n">data</span><span class="o">.</span><span class="n">flatten</span><span class="p">(</span><span class="n">order</span><span class="o">=</span><span class="s1">'C'</span><span class="p">)</span>
<span class="n">header</span> <span class="o">=</span> <span class="s2">"""# OpenDX density file</span>
<span class="s2"># File format: http://opendx.sdsc.edu/docs/html/pages/usrgu068.htm#HDREDF</span>
<span class="s2">object 1 class gridpositions counts </span><span class="si">{nx}</span><span class="s2"> </span><span class="si">{ny}</span><span class="s2"> </span><span class="si">{nz}</span>
<span class="s2">origin </span><span class="si">{ox}</span><span class="s2"> </span><span class="si">{oy}</span><span class="s2"> </span><span class="si">{oz}</span>
<span class="s2">delta </span><span class="si">{dx}</span><span class="s2"> 0.000000 0.000000</span>
<span class="s2">delta 0.000000 </span><span class="si">{dy}</span><span class="s2"> 0.000000</span>
<span class="s2">delta 0.000000 0.000000 </span><span class="si">{dz}</span>
<span class="s2">object 2 class gridconnections counts </span><span class="si">{nx}</span><span class="s2"> </span><span class="si">{ny}</span><span class="s2"> </span><span class="si">{nz}</span>
<span class="s2">object 3 class array type double rank 0 items </span><span class="si">{num}</span><span class="s2"> data follows"""</span><span class="o">.</span><span class="n">format</span><span class="p">(</span><span class="o">**</span><span class="n">headerInfo</span><span class="p">)</span>
<span class="nb">len</span><span class="p">(</span><span class="n">data</span><span class="p">)</span>
<span class="k">if</span> <span class="n">num</span> <span class="o">==</span> <span class="mi">3</span><span class="o">*</span><span class="p">(</span><span class="n">num</span><span class="o">//</span><span class="mi">3</span><span class="p">):</span>
<span class="n">footer</span> <span class="o">=</span> <span class="s2">""</span>
<span class="k">else</span><span class="p">:</span>
<span class="n">footer</span> <span class="o">=</span> <span class="s2">" "</span><span class="o">.</span><span class="n">join</span><span class="p">([</span><span class="n">fmt</span> <span class="o">%</span> <span class="n">x</span> <span class="k">for</span> <span class="n">x</span> <span class="ow">in</span> <span class="n">data</span><span class="p">[</span><span class="mi">3</span><span class="o">*</span><span class="p">(</span><span class="n">num</span><span class="o">//</span><span class="mi">3</span><span class="p">):]])</span> <span class="c1"># last line of data</span>
<span class="n">footer</span> <span class="o">+=</span> <span class="s2">"</span><span class="se">\n</span><span class="s2">"</span>
<span class="n">footer</span> <span class="o">+=</span> <span class="s2">"""attribute "dep" string "positions"</span>
<span class="s2">object "density" class field </span>
<span class="s2">component "positions" value 1</span>
<span class="s2">component "connections" value 2</span>
<span class="s2">component "data" value 3</span>
<span class="s2">"""</span>
<span class="n">np</span><span class="o">.</span><span class="n">savetxt</span><span class="p">(</span> <span class="n">outfile</span><span class="p">,</span> <span class="n">np</span><span class="o">.</span><span class="n">reshape</span><span class="p">(</span><span class="n">data</span><span class="p">[:</span><span class="mi">3</span><span class="o">*</span><span class="p">(</span><span class="n">num</span><span class="o">//</span><span class="mi">3</span><span class="p">)],</span> <span class="p">(</span><span class="n">num</span><span class="o">//</span><span class="mi">3</span><span class="p">,</span><span class="mi">3</span><span class="p">),</span> <span class="n">order</span><span class="o">=</span><span class="s1">'C'</span><span class="p">),</span>
<span class="n">fmt</span><span class="o">=</span><span class="n">fmt</span><span class="p">,</span>
<span class="n">header</span><span class="o">=</span><span class="n">header</span><span class="p">,</span> <span class="n">comments</span><span class="o">=</span><span class="s1">''</span><span class="p">,</span> <span class="n">footer</span><span class="o">=</span><span class="n">footer</span> <span class="p">)</span></div>
<div class="viewcode-block" id="add_smaller_grid">
<a class="viewcode-back" href="../../api/utilities/grid.html#arbdmodel.grid.add_smaller_grid">[docs]</a>
<span class="k">def</span><span class="w"> </span><span class="nf">add_smaller_grid</span><span class="p">(</span> <span class="n">grid</span><span class="p">,</span> <span class="n">smaller_grid</span> <span class="p">):</span>
<span class="n">slices</span> <span class="o">=</span> <span class="n">get_slice_enclosing_smaller_grid</span><span class="p">(</span> <span class="n">grid</span><span class="p">,</span> <span class="n">smaller_grid</span> <span class="p">)</span>
<span class="n">grid</span><span class="o">.</span><span class="n">grid</span><span class="p">[</span><span class="n">slices</span><span class="p">]</span> <span class="o">=</span> <span class="n">grid</span><span class="o">.</span><span class="n">grid</span><span class="p">[</span><span class="n">slices</span><span class="p">]</span> <span class="o">+</span> <span class="n">smaller_grid</span><span class="o">.</span><span class="n">grid</span></div>
<div class="viewcode-block" id="average_grids">
<a class="viewcode-back" href="../../api/utilities/grid.html#arbdmodel.grid.average_grids">[docs]</a>
<span class="k">def</span><span class="w"> </span><span class="nf">average_grids</span><span class="p">(</span><span class="n">grids</span><span class="p">,</span> <span class="n">mask</span><span class="o">=</span><span class="s1">'nan'</span><span class="p">):</span>
<span class="w"> </span><span class="sd">"""</span>
<span class="sd"> Compute the average of multiple grids.</span>
<span class="sd"> Parameters</span>
<span class="sd"> ----------</span>
<span class="sd"> param grids: The input grids to average.</span>
<span class="sd"> type grids: list of ndarrays</span>
<span class="sd"> param mask: The mask type to use for averaging. If 'nan', only non-NaN values are considered. Default is 'nan'.</span>
<span class="sd"> type mask: str</span>
<span class="sd"> return: The average grid.</span>
<span class="sd"> rtype: ndarray</span>
<span class="sd"> raises NotImplementedError: If the mask option is not implemented.</span>
<span class="sd"> Example:</span>
<span class="sd"> ----------</span>
<span class="sd"> >>> grid1 = np.array([[1, 2, 3], [4, 5, 6], [7, 8, 9]])</span>
<span class="sd"> >>> grid2 = np.array([[2, 4, 6], [8, 10, 12], [14, 16, 18]])</span>
<span class="sd"> >>> averaged_grid = average_grids([grid1, grid2], mask='nan')</span>
<span class="sd"> >>> print(averaged_grid)</span>
<span class="sd"> """</span>
<span class="k">def</span><span class="w"> </span><span class="nf">__get_mask</span><span class="p">(</span><span class="n">grid</span><span class="p">,</span> <span class="n">mask</span><span class="p">):</span> <span class="k">if</span> <span class="n">mask</span> <span class="o">==</span> <span class="s1">'nan'</span><span class="p">:</span>
<span class="n">g_mask</span> <span class="o">=</span> <span class="o">~</span><span class="n">np</span><span class="o">.</span><span class="n">isnan</span><span class="p">(</span><span class="n">grid</span><span class="p">)</span>
<span class="k">else</span><span class="p">:</span>
<span class="k">raise</span> <span class="ne">NotImplementedError</span>
<span class="k">return</span> <span class="n">g_mask</span>
<span class="n">counts</span> <span class="o">=</span> <span class="n">np</span><span class="o">.</span><span class="n">zeros</span><span class="p">(</span><span class="n">grids</span><span class="p">[</span><span class="mi">0</span><span class="p">]</span><span class="o">.</span><span class="n">shape</span><span class="p">,</span> <span class="n">dtype</span><span class="o">=</span><span class="nb">int</span><span class="p">)</span>
<span class="n">totals</span> <span class="o">=</span> <span class="n">np</span><span class="o">.</span><span class="n">zeros</span><span class="p">(</span><span class="n">counts</span><span class="o">.</span><span class="n">shape</span><span class="p">)</span>
<span class="k">for</span> <span class="n">g</span> <span class="ow">in</span> <span class="n">grids</span><span class="p">:</span>
<span class="n">g_mask</span> <span class="o">=</span> <span class="n">__get_mask</span><span class="p">(</span><span class="n">g</span><span class="p">,</span> <span class="n">mask</span><span class="p">)</span>
<span class="n">counts</span> <span class="o">=</span> <span class="n">counts</span> <span class="o">+</span> <span class="n">g_mask</span>
<span class="n">totals</span><span class="p">[</span><span class="n">g_mask</span><span class="p">]</span> <span class="o">=</span> <span class="n">totals</span><span class="p">[</span><span class="n">g_mask</span><span class="p">]</span> <span class="o">+</span> <span class="n">g</span><span class="p">[</span><span class="n">g_mask</span><span class="p">]</span>
<span class="n">sl</span> <span class="o">=</span> <span class="n">counts</span> <span class="o">></span> <span class="mi">0</span>
<span class="n">average</span> <span class="o">=</span> <span class="n">totals</span>
<span class="n">average</span><span class="p">[</span><span class="n">sl</span><span class="p">]</span> <span class="o">=</span> <span class="n">totals</span><span class="p">[</span><span class="n">sl</span><span class="p">]</span> <span class="o">/</span> <span class="n">counts</span><span class="p">[</span><span class="n">sl</span><span class="p">]</span>
<span class="n">average</span><span class="p">[</span><span class="n">counts</span> <span class="o">==</span> <span class="mi">0</span><span class="p">]</span> <span class="o">=</span> <span class="n">np</span><span class="o">.</span><span class="n">nan</span>
<span class="k">return</span> <span class="n">average</span></div>
<div class="viewcode-block" id="loadGrid">
<a class="viewcode-back" href="../../api/utilities/grid.html#arbdmodel.grid.loadGrid">[docs]</a>
<span class="k">def</span><span class="w"> </span><span class="nf">loadGrid</span><span class="p">(</span><span class="n">file</span><span class="p">,</span> <span class="o">**</span><span class="n">kwargs</span><span class="p">):</span>
<span class="w"> </span><span class="sd">"""Load grid data from DX file"""</span>
<span class="c1"># Read DX file header to get dimensions and metadata</span>
<span class="k">with</span> <span class="nb">open</span><span class="p">(</span><span class="n">file</span><span class="p">,</span> <span class="s1">'r'</span><span class="p">)</span> <span class="k">as</span> <span class="n">f</span><span class="p">:</span>
<span class="n">lines</span> <span class="o">=</span> <span class="n">f</span><span class="o">.</span><span class="n">readlines</span><span class="p">()</span>
<span class="c1"># Parse header</span>
<span class="k">for</span> <span class="n">line</span> <span class="ow">in</span> <span class="n">lines</span><span class="p">:</span>
<span class="k">if</span> <span class="s1">'counts'</span> <span class="ow">in</span> <span class="n">line</span><span class="p">:</span>
<span class="n">dims</span> <span class="o">=</span> <span class="p">[</span><span class="nb">int</span><span class="p">(</span><span class="n">x</span><span class="p">)</span> <span class="k">for</span> <span class="n">x</span> <span class="ow">in</span> <span class="n">line</span><span class="o">.</span><span class="n">split</span><span class="p">()[</span><span class="o">-</span><span class="mi">3</span><span class="p">:]]</span>
<span class="k">break</span>
<span class="k">for</span> <span class="n">line</span> <span class="ow">in</span> <span class="n">lines</span><span class="p">:</span>
<span class="k">if</span> <span class="s1">'origin'</span> <span class="ow">in</span> <span class="n">line</span><span class="p">:</span>
<span class="n">origin</span> <span class="o">=</span> <span class="p">[</span><span class="nb">float</span><span class="p">(</span><span class="n">x</span><span class="p">)</span> <span class="k">for</span> <span class="n">x</span> <span class="ow">in</span> <span class="n">line</span><span class="o">.</span><span class="n">split</span><span class="p">()[</span><span class="o">-</span><span class="mi">3</span><span class="p">:]]</span>
<span class="k">break</span>
<span class="k">for</span> <span class="n">line</span> <span class="ow">in</span> <span class="n">lines</span><span class="p">:</span>
<span class="k">if</span> <span class="s1">'delta'</span> <span class="ow">in</span> <span class="n">line</span><span class="p">:</span>
<span class="n">delta</span> <span class="o">=</span> <span class="nb">float</span><span class="p">(</span><span class="n">line</span><span class="o">.</span><span class="n">split</span><span class="p">()[</span><span class="mi">1</span><span class="p">])</span>
<span class="k">break</span>
<span class="c1"># Read data values</span>
<span class="n">data</span> <span class="o">=</span> <span class="p">[]</span>
<span class="k">for</span> <span class="n">line</span> <span class="ow">in</span> <span class="n">lines</span><span class="p">:</span>
<span class="k">try</span><span class="p">:</span>
<span class="n">values</span> <span class="o">=</span> <span class="p">[</span><span class="nb">float</span><span class="p">(</span><span class="n">x</span><span class="p">)</span> <span class="k">for</span> <span class="n">x</span> <span class="ow">in</span> <span class="n">line</span><span class="o">.</span><span class="n">split</span><span class="p">()]</span>
<span class="n">data</span><span class="o">.</span><span class="n">extend</span><span class="p">(</span><span class="n">values</span><span class="p">)</span>
<span class="k">except</span> <span class="ne">ValueError</span><span class="p">:</span>
<span class="k">continue</span>
<span class="c1"># Reshape into 3D array</span>
<span class="n">grid</span> <span class="o">=</span> <span class="n">np</span><span class="o">.</span><span class="n">array</span><span class="p">(</span><span class="n">data</span><span class="p">)</span><span class="o">.</span><span class="n">reshape</span><span class="p">(</span><span class="n">dims</span><span class="p">)</span>
<span class="k">return</span> <span class="n">grid</span><span class="p">,</span> <span class="n">origin</span><span class="p">,</span> <span class="n">delta</span></div>
<div class="viewcode-block" id="Bound_grid">
<a class="viewcode-back" href="../../api/utilities/grid.html#arbdmodel.grid.Bound_grid">[docs]</a>
<span class="k">def</span><span class="w"> </span><span class="nf">Bound_grid</span><span class="p">(</span><span class="n">inFile</span><span class="p">,</span> <span class="n">outFile</span><span class="p">,</span> <span class="n">lowerBound</span><span class="p">,</span> <span class="n">upperBound</span><span class="p">):</span>
<span class="w"> </span><span class="sd">"""Apply bounds to grid values"""</span>
<span class="c1"># Fix scientific notation</span>
<span class="n">cmd_in</span> <span class="o">=</span> <span class="s2">"sed -r 's/^([0-9]+)e/</span><span class="se">\1</span><span class="s2">.0e/g; s/ ([0-9]+)e/ </span><span class="se">\1</span><span class="s2">.0e/' "</span> <span class="o">+</span> <span class="n">inFile</span> <span class="o">+</span> <span class="s2">" > bound_grid_temp0.dx"</span>
<span class="n">os</span><span class="o">.</span><span class="n">system</span><span class="p">(</span><span class="n">cmd_in</span><span class="p">)</span>
<span class="n">cmd_in</span> <span class="o">=</span> <span class="s2">"sed -r 's/^(-[0-9]+)e/</span><span class="se">\1</span><span class="s2">.0e/g; s/ (-[0-9]+)e/ </span><span class="se">\1</span><span class="s2">.0e/' bound_grid_temp0.dx > bound_grid_temp1.dx"</span>
<span class="n">os</span><span class="o">.</span><span class="n">system</span><span class="p">(</span><span class="n">cmd_in</span><span class="p">)</span>
<span class="k">assert</span><span class="p">(</span><span class="n">lowerBound</span> <span class="o"><</span> <span class="n">upperBound</span><span class="p">)</span>
<span class="c1"># Load data</span>
<span class="n">grid</span><span class="p">,</span> <span class="n">origin</span><span class="p">,</span> <span class="n">delta</span> <span class="o">=</span> <span class="n">loadGrid</span><span class="p">(</span><span class="s1">'bound_grid_temp1.dx'</span><span class="p">)</span>
<span class="c1"># Apply bounds</span>
<span class="n">grid</span><span class="p">[</span><span class="n">grid</span> <span class="o">></span> <span class="n">upperBound</span><span class="p">]</span> <span class="o">=</span> <span class="n">upperBound</span>
<span class="n">grid</span><span class="p">[</span><span class="n">grid</span> <span class="o"><</span> <span class="n">lowerBound</span><span class="p">]</span> <span class="o">=</span> <span class="n">lowerBound</span>
<span class="c1"># Write output</span>
<span class="n">writeDx</span><span class="p">(</span><span class="n">outFile</span><span class="p">,</span> <span class="n">grid</span><span class="p">,</span> <span class="n">origin</span><span class="p">,</span> <span class="p">[</span><span class="n">delta</span><span class="p">,</span> <span class="n">delta</span><span class="p">,</span> <span class="n">delta</span><span class="p">])</span></div>
<div class="viewcode-block" id="Create_null">
<a class="viewcode-back" href="../../api/utilities/grid.html#arbdmodel.grid.Create_null">[docs]</a>
<span class="k">def</span><span class="w"> </span><span class="nf">Create_null</span><span class="p">(</span><span class="n">grid_path</span><span class="o">=</span><span class="s1">'null.dx'</span><span class="p">):</span>
<span class="w"> </span><span class="sd">"""Create null potential grid"""</span>
<span class="n">zeros</span> <span class="o">=</span> <span class="n">np</span><span class="o">.</span><span class="n">zeros</span><span class="p">([</span><span class="mi">2</span><span class="p">,</span><span class="mi">2</span><span class="p">,</span><span class="mi">2</span><span class="p">])</span>
<span class="n">origin</span> <span class="o">=</span> <span class="o">-</span><span class="mi">1500</span><span class="o">*</span><span class="n">np</span><span class="o">.</span><span class="n">array</span><span class="p">((</span><span class="mi">1</span><span class="p">,</span><span class="mi">1</span><span class="p">,</span><span class="mi">1</span><span class="p">))</span>
<span class="n">delta</span> <span class="o">=</span> <span class="p">[</span><span class="mi">3000</span><span class="p">,</span> <span class="mi">3000</span><span class="p">,</span> <span class="mi">3000</span><span class="p">]</span>
<span class="n">writeDx</span><span class="p">(</span><span class="n">grid_path</span><span class="p">,</span> <span class="n">zeros</span><span class="p">,</span> <span class="n">origin</span><span class="p">,</span> <span class="n">delta</span><span class="p">)</span></div>
<div class="viewcode-block" id="TestAverageGrids">
<a class="viewcode-back" href="../../api/utilities/grid.html#arbdmodel.grid.TestAverageGrids">[docs]</a>
<span class="k">class</span><span class="w"> </span><span class="nc">TestAverageGrids</span><span class="p">(</span><span class="n">unittest</span><span class="o">.</span><span class="n">TestCase</span><span class="p">):</span>
<div class="viewcode-block" id="TestAverageGrids.test_average_grids">
<a class="viewcode-back" href="../../api/utilities/grid.html#arbdmodel.grid.TestAverageGrids.test_average_grids">[docs]</a>
<span class="k">def</span><span class="w"> </span><span class="nf">test_average_grids</span><span class="p">(</span><span class="bp">self</span><span class="p">):</span>
<span class="n">grid1</span> <span class="o">=</span> <span class="n">np</span><span class="o">.</span><span class="n">array</span><span class="p">([[</span><span class="mi">1</span><span class="p">,</span> <span class="mi">2</span><span class="p">,</span> <span class="n">np</span><span class="o">.</span><span class="n">nan</span><span class="p">],</span> <span class="p">[</span><span class="mi">4</span><span class="p">,</span> <span class="mi">5</span><span class="p">,</span> <span class="mi">6</span><span class="p">],</span> <span class="p">[</span><span class="mi">7</span><span class="p">,</span> <span class="mi">8</span><span class="p">,</span> <span class="mi">9</span><span class="p">]])</span>
<span class="n">grid2</span> <span class="o">=</span> <span class="n">np</span><span class="o">.</span><span class="n">array</span><span class="p">([[</span><span class="mi">2</span><span class="p">,</span> <span class="mi">4</span><span class="p">,</span> <span class="mi">6</span><span class="p">],</span> <span class="p">[</span><span class="mi">8</span><span class="p">,</span> <span class="n">np</span><span class="o">.</span><span class="n">nan</span><span class="p">,</span> <span class="mi">12</span><span class="p">],</span> <span class="p">[</span><span class="mi">14</span><span class="p">,</span> <span class="mi">16</span><span class="p">,</span> <span class="mi">18</span><span class="p">]])</span>
<span class="n">expected_result</span> <span class="o">=</span> <span class="n">np</span><span class="o">.</span><span class="n">array</span><span class="p">([[</span><span class="mf">1.5</span><span class="p">,</span> <span class="mi">3</span><span class="p">,</span> <span class="mi">6</span><span class="p">],</span> <span class="p">[</span><span class="mi">6</span><span class="p">,</span> <span class="mi">5</span><span class="p">,</span> <span class="mi">9</span><span class="p">],</span> <span class="p">[</span><span class="mf">10.5</span><span class="p">,</span> <span class="mi">12</span><span class="p">,</span> <span class="mf">13.5</span><span class="p">]])</span>
<span class="n">result</span> <span class="o">=</span> <span class="n">average_grids</span><span class="p">([</span><span class="n">grid1</span><span class="p">,</span> <span class="n">grid2</span><span class="p">],</span> <span class="n">mask</span><span class="o">=</span><span class="s1">'nan'</span><span class="p">)</span>
<span class="n">np</span><span class="o">.</span><span class="n">testing</span><span class="o">.</span><span class="n">assert_array_equal</span><span class="p">(</span><span class="n">result</span><span class="p">,</span> <span class="n">expected_result</span><span class="p">)</span></div>
</div>
<div class="viewcode-block" id="neighborhood_average">
<a class="viewcode-back" href="../../api/utilities/grid.html#arbdmodel.grid.neighborhood_average">[docs]</a>
<span class="k">def</span><span class="w"> </span><span class="nf">neighborhood_average</span><span class="p">(</span><span class="n">grid</span><span class="p">,</span> <span class="n">neighborhood</span> <span class="o">=</span> <span class="mi">1</span><span class="p">,</span> <span class="n">fill_value</span><span class="o">=</span><span class="s1">'mirror'</span><span class="p">):</span>
<span class="w"> </span><span class="sd">"""</span>
<span class="sd"> Compute the neighborhood average of a grid.</span>
<span class="sd"> Parameters:</span>
<span class="sd"> grid (ndarray): The input grid.</span>
<span class="sd"> neighborhood (int or list): The size of the neighborhood to use for averaging. If an integer is provided, the same neighborhood size is used for all dimensions. If a list is provided, it should specify the neighborhood size for each dimension separately. Default is 1.</span>
<span class="sd"> fill_value (str or ndarray): The fill value to use for padding the grid. If 'mirror', values are mirrored along each dimension. If 'nearest', values are filled with the nearest neighbor. If 'periodic', values are filled with periodic boundary conditions. If a number is provided, it is used directly as the fill value. Default is 'mirror'.</span>
<span class="sd"> Returns:</span>
<span class="sd"> ndarray: The grid with the neighborhood average computed.</span>
<span class="sd"> Note:</span>
<span class="sd"> - Currently, only the 'mirror' fill value option is implemented.</span>
<span class="sd"> - This function requires the 'average_grids' function from an external source.</span>
<span class="sd"> Example:</span>
<span class="sd"> >>> grid = np.array([[1, 2, 3], [4, 5, 6], [7, 8, 9]])</span>
<span class="sd"> >>> averaged_grid = neighborhood_average(grid, neighborhood=2, fill_value='mirror')</span>
<span class="sd"> >>> print(averaged_grid)</span>
<span class="sd"> """</span>
<span class="k">try</span><span class="p">:</span>
<span class="n">neighborhood</span><span class="p">[</span><span class="mi">0</span><span class="p">]</span>
<span class="k">except</span><span class="p">:</span>
<span class="n">neighborhood</span> <span class="o">=</span> <span class="p">[</span><span class="n">neighborhood</span><span class="p">]</span> <span class="o">*</span> <span class="n">grid</span><span class="o">.</span><span class="n">ndim</span>
<span class="n">padded_grid</span> <span class="o">=</span> <span class="n">np</span><span class="o">.</span><span class="n">ones</span><span class="p">(</span> <span class="p">[</span><span class="n">s</span><span class="o">+</span><span class="n">dx</span><span class="o">*</span><span class="mi">2</span> <span class="k">for</span> <span class="n">s</span><span class="p">,</span><span class="n">dx</span> <span class="ow">in</span> <span class="nb">zip</span><span class="p">(</span><span class="n">grid</span><span class="o">.</span><span class="n">shape</span><span class="p">,</span><span class="n">neighborhood</span><span class="p">)]</span> <span class="p">)</span><span class="o">*</span><span class="n">np</span><span class="o">.</span><span class="n">nan</span>
<span class="k">if</span> <span class="n">fill_value</span> <span class="o">==</span> <span class="s1">'mirror'</span><span class="p">:</span> <span class="k">def</span><span class="w"> </span><span class="nf">get_slices</span><span class="p">(</span><span class="n">mirror_dims</span><span class="p">,</span> <span class="n">left</span><span class="p">):</span>
<span class="n">sl1</span> <span class="o">=</span> <span class="nb">tuple</span><span class="p">(</span><span class="nb">slice</span><span class="p">(</span><span class="n">n</span><span class="p">,</span><span class="o">-</span><span class="n">n</span><span class="p">)</span> <span class="k">if</span> <span class="n">i</span> <span class="ow">not</span> <span class="ow">in</span> <span class="n">mirror_dims</span> <span class="k">else</span>
<span class="nb">slice</span><span class="p">(</span><span class="kc">None</span><span class="p">,</span><span class="n">n</span><span class="p">)</span> <span class="k">if</span> <span class="n">left</span><span class="p">[</span><span class="n">mirror_dims</span><span class="o">.</span><span class="n">index</span><span class="p">(</span><span class="n">i</span><span class="p">)]</span>
<span class="k">else</span> <span class="nb">slice</span><span class="p">(</span><span class="o">-</span><span class="n">n</span><span class="p">,</span><span class="kc">None</span><span class="p">)</span>
<span class="k">for</span> <span class="n">i</span><span class="p">,</span><span class="n">n</span> <span class="ow">in</span> <span class="nb">enumerate</span><span class="p">(</span><span class="n">neighborhood</span><span class="p">))</span>
<span class="n">sl2</span> <span class="o">=</span> <span class="nb">tuple</span><span class="p">(</span><span class="nb">slice</span><span class="p">(</span><span class="kc">None</span><span class="p">,</span><span class="kc">None</span><span class="p">)</span> <span class="k">if</span> <span class="n">i</span> <span class="ow">not</span> <span class="ow">in</span> <span class="n">mirror_dims</span> <span class="k">else</span>
<span class="nb">slice</span><span class="p">((</span><span class="n">n</span><span class="o">-</span><span class="mi">1</span><span class="p">),</span><span class="kc">None</span><span class="p">,</span><span class="o">-</span><span class="mi">1</span><span class="p">)</span> <span class="k">if</span> <span class="n">left</span><span class="p">[</span><span class="n">mirror_dims</span><span class="o">.</span><span class="n">index</span><span class="p">(</span><span class="n">i</span><span class="p">)]</span>
<span class="k">else</span> <span class="nb">slice</span><span class="p">(</span><span class="kc">None</span><span class="p">,</span><span class="o">-</span><span class="p">(</span><span class="n">n</span><span class="o">+</span><span class="mi">1</span><span class="p">),</span><span class="o">-</span><span class="mi">1</span><span class="p">)</span>
<span class="k">for</span> <span class="n">i</span><span class="p">,</span><span class="n">n</span> <span class="ow">in</span> <span class="nb">enumerate</span><span class="p">(</span><span class="n">neighborhood</span><span class="p">))</span>
<span class="k">return</span> <span class="n">sl1</span><span class="p">,</span><span class="n">sl2</span>
<span class="k">def</span><span class="w"> </span><span class="nf">build_dim_arrays</span><span class="p">(</span><span class="n">nd</span><span class="p">,</span> <span class="n">mirror_dims</span><span class="p">,</span> <span class="n">left</span><span class="p">):</span>
<span class="n">d0</span> <span class="o">=</span> <span class="n">mirror_dims</span><span class="p">[</span><span class="o">-</span><span class="mi">1</span><span class="p">]</span> <span class="k">if</span> <span class="nb">len</span><span class="p">(</span><span class="n">mirror_dims</span><span class="p">)</span> <span class="o">></span> <span class="mi">0</span> <span class="k">else</span> <span class="o">-</span><span class="mi">1</span>
<span class="n">mirror_dims</span><span class="o">.</span><span class="n">append</span><span class="p">(</span><span class="n">d0</span><span class="p">)</span>
<span class="n">left</span><span class="o">.</span><span class="n">append</span><span class="p">(</span><span class="kc">False</span><span class="p">)</span>
<span class="k">for</span> <span class="n">d</span> <span class="ow">in</span> <span class="nb">range</span><span class="p">(</span><span class="n">d0</span><span class="o">+</span><span class="mi">1</span><span class="p">,</span><span class="n">grid</span><span class="o">.</span><span class="n">ndim</span><span class="o">+</span><span class="mi">1</span><span class="o">-</span><span class="n">nd</span><span class="p">):</span>
<span class="n">mirror_dims</span><span class="p">[</span><span class="o">-</span><span class="mi">1</span><span class="p">]</span> <span class="o">=</span> <span class="n">d</span>
<span class="n">left</span><span class="p">[</span><span class="o">-</span><span class="mi">1</span><span class="p">]</span> <span class="o">=</span> <span class="kc">False</span>
<span class="k">yield</span> <span class="nb">list</span><span class="p">(</span><span class="n">mirror_dims</span><span class="p">),</span> <span class="nb">list</span><span class="p">(</span><span class="n">left</span><span class="p">)</span>
<span class="n">left</span><span class="p">[</span><span class="o">-</span><span class="mi">1</span><span class="p">]</span> <span class="o">=</span> <span class="kc">True</span>
<span class="k">yield</span> <span class="nb">list</span><span class="p">(</span><span class="n">mirror_dims</span><span class="p">),</span> <span class="nb">list</span><span class="p">(</span><span class="n">left</span><span class="p">)</span>
<span class="n">all_mirror_dims_arrays</span> <span class="o">=</span> <span class="p">[]</span>
<span class="n">all_left_arrays</span> <span class="o">=</span> <span class="p">[]</span>
<span class="k">for</span> <span class="n">nd</span> <span class="ow">in</span> <span class="nb">range</span><span class="p">(</span><span class="mi">1</span><span class="p">,</span><span class="n">grid</span><span class="o">.</span><span class="n">ndim</span><span class="o">+</span><span class="mi">1</span><span class="p">):</span>
<span class="c1"># print(nd,grid.ndim)</span>
<span class="n">mirror_dims_arrays</span> <span class="o">=</span> <span class="p">[[]]</span>
<span class="n">left_arrays</span> <span class="o">=</span> <span class="p">[[]]</span>
<span class="k">while</span> <span class="n">nd</span> <span class="o">></span> <span class="mi">0</span><span class="p">:</span>
<span class="k">for</span> <span class="n">m0</span><span class="p">,</span><span class="n">l0</span> <span class="ow">in</span> <span class="nb">list</span><span class="p">(</span><span class="nb">zip</span><span class="p">(</span><span class="n">mirror_dims_arrays</span><span class="p">,</span><span class="n">left_arrays</span><span class="p">)):</span>
<span class="k">for</span> <span class="n">m</span><span class="p">,</span><span class="n">l</span> <span class="ow">in</span> <span class="n">build_dim_arrays</span><span class="p">(</span><span class="n">nd</span><span class="p">,</span> <span class="n">m0</span><span class="p">,</span> <span class="n">l0</span><span class="p">):</span>
<span class="n">mirror_dims_arrays</span><span class="o">.</span><span class="n">append</span><span class="p">(</span><span class="n">m</span><span class="p">)</span>
<span class="n">left_arrays</span><span class="o">.</span><span class="n">append</span><span class="p">(</span><span class="n">l</span><span class="p">)</span>
<span class="n">nd</span><span class="o">=</span><span class="n">nd</span><span class="o">-</span><span class="mi">1</span>
<span class="c1"># print("dbg",nd) </span>
<span class="c1"># print(mirror_dims_arrays)</span>
<span class="n">all_mirror_dims_arrays</span><span class="o">.</span><span class="n">extend</span><span class="p">(</span><span class="n">mirror_dims_arrays</span><span class="p">)</span>
<span class="n">all_left_arrays</span><span class="o">.</span><span class="n">extend</span><span class="p">(</span><span class="n">left_arrays</span><span class="p">)</span>
<span class="k">for</span> <span class="n">m</span><span class="p">,</span><span class="n">l</span> <span class="ow">in</span> <span class="nb">zip</span><span class="p">(</span><span class="n">all_mirror_dims_arrays</span><span class="p">,</span> <span class="n">all_left_arrays</span><span class="p">):</span>
<span class="c1">## set values mirroring d</span>
<span class="n">sl1</span><span class="p">,</span><span class="n">sl2</span> <span class="o">=</span> <span class="n">get_slices</span><span class="p">(</span><span class="n">m</span><span class="p">,</span> <span class="n">l</span><span class="p">)</span>
<span class="n">padded_grid</span><span class="p">[</span><span class="n">sl1</span><span class="p">]</span> <span class="o">=</span> <span class="n">grid</span><span class="p">[</span><span class="n">sl2</span><span class="p">]</span>
<span class="k">elif</span> <span class="n">fill_value</span> <span class="o">==</span> <span class="s1">'nearest'</span><span class="p">:</span>
<span class="k">raise</span> <span class="ne">NotImplementedError</span>
<span class="k">elif</span> <span class="n">fill_value</span> <span class="o">==</span> <span class="s1">'periodic'</span><span class="p">:</span>
<span class="k">raise</span> <span class="ne">NotImplementedError</span>
<span class="k">else</span><span class="p">:</span>
<span class="n">padded_grid</span> <span class="o">=</span> <span class="n">fill_value</span>
<span class="n">sl</span> <span class="o">=</span> <span class="nb">tuple</span><span class="p">(</span><span class="nb">slice</span><span class="p">(</span><span class="n">n</span><span class="p">,</span><span class="o">-</span><span class="n">n</span><span class="p">)</span> <span class="k">for</span> <span class="n">n</span> <span class="ow">in</span> <span class="n">neighborhood</span><span class="p">)</span>
<span class="n">padded_grid</span><span class="p">[</span><span class="n">sl</span><span class="p">]</span> <span class="o">=</span> <span class="n">grid</span>
<span class="c1"># print(padded_grid)</span>
<span class="c1">## Gather sliced arrays to perform average</span>
<span class="k">def</span><span class="w"> </span><span class="nf">get_slice</span><span class="p">(</span><span class="n">neighborhood</span><span class="p">,</span> <span class="n">slices</span><span class="o">=</span><span class="kc">None</span><span class="p">):</span>
<span class="c1"># print("get_slice",neighborhood)</span>
<span class="k">if</span> <span class="n">slices</span> <span class="ow">is</span> <span class="kc">None</span><span class="p">:</span> <span class="n">slices</span> <span class="o">=</span> <span class="p">[[]]</span>
<span class="k">if</span> <span class="nb">len</span><span class="p">(</span><span class="n">neighborhood</span><span class="p">)</span> <span class="o">></span> <span class="mi">1</span><span class="p">:</span>
<span class="n">slices</span> <span class="o">=</span> <span class="n">get_slice</span><span class="p">(</span><span class="n">neighborhood</span><span class="p">[</span><span class="mi">1</span><span class="p">:],</span> <span class="n">slices</span><span class="p">)</span>
<span class="n">nmax</span> <span class="o">=</span> <span class="mi">2</span><span class="o">*</span><span class="n">neighborhood</span><span class="p">[</span><span class="mi">0</span><span class="p">]</span>
<span class="n">slices</span> <span class="o">=</span> <span class="p">[[</span><span class="nb">slice</span><span class="p">(</span><span class="n">n</span><span class="p">,</span><span class="o">-</span><span class="p">(</span><span class="n">nmax</span><span class="o">-</span><span class="n">n</span><span class="p">)</span> <span class="k">if</span> <span class="n">nmax</span> <span class="o">!=</span> <span class="n">n</span> <span class="k">else</span> <span class="kc">None</span><span class="p">)]</span><span class="o">+</span><span class="n">sls</span> <span class="k">for</span> <span class="n">n</span> <span class="ow">in</span> <span class="nb">range</span><span class="p">(</span><span class="n">nmax</span><span class="o">+</span><span class="mi">1</span><span class="p">)</span> <span class="k">for</span> <span class="n">sls</span> <span class="ow">in</span> <span class="n">slices</span><span class="p">]</span>
<span class="k">return</span> <span class="n">slices</span>
<span class="n">slices</span> <span class="o">=</span> <span class="n">get_slice</span><span class="p">(</span><span class="n">neighborhood</span><span class="p">)</span>
<span class="n">result</span> <span class="o">=</span> <span class="n">average_grids</span><span class="p">([</span><span class="n">padded_grid</span><span class="p">[</span><span class="n">sl</span><span class="p">]</span> <span class="k">for</span> <span class="n">sl</span> <span class="ow">in</span> <span class="n">slices</span><span class="p">])</span>
<span class="k">return</span> <span class="n">result</span></div>
<div class="viewcode-block" id="replace_false_with_distance"><a class="viewcode-back" href="../../api/utilities/grid.html#arbdmodel.grid.replace_false_with_distance">[docs]</a>
<span class="k">def</span><span class="w"> </span><span class="nf">replace_false_with_distance</span><span class="p">(</span><span class="n">boolean_grid</span><span class="p">,</span> <span class="n">sampling</span><span class="p">):</span>
<span class="kn">import</span><span class="w"> </span><span class="nn">ndimage</span>
<span class="k">return</span> <span class="n">ndimage</span><span class="o">.</span><span class="n">morphology</span><span class="o">.</span><span class="n">distance_transform_edt</span><span class="p">(</span> <span class="n">boolean_grid</span><span class="p">,</span> <span class="n">sampling</span><span class="o">=</span><span class="p">[</span><span class="n">Dxy</span><span class="p">,</span><span class="n">Dxy</span><span class="p">,</span><span class="n">Dz</span><span class="p">]</span> <span class="p">)</span></div>
<div class="viewcode-block" id="fill_nans">
<a class="viewcode-back" href="../../api/utilities/grid.html#arbdmodel.grid.fill_nans">[docs]</a>
<span class="k">def</span><span class="w"> </span><span class="nf">fill_nans</span><span class="p">(</span><span class="n">grid</span><span class="p">,</span> <span class="n">neighborhood</span><span class="o">=</span><span class="mi">1</span><span class="p">,</span> <span class="n">max_iterations</span><span class="o">=</span><span class="n">np</span><span class="o">.</span><span class="n">inf</span><span class="p">,</span> <span class="n">mask</span><span class="o">=</span><span class="kc">None</span><span class="p">):</span>
<span class="w"> </span><span class="sd">"""</span>
<span class="sd"> Fill NaN values in a grid using neighborhood averaging.</span>
<span class="sd"> Parameters:</span>
<span class="sd"> grid (ndarray): The input grid containing NaN values.</span>
<span class="sd"> neighborhood (int, optional): The size of the neighborhood to use for averaging. Default is 1.</span>
<span class="sd"> max_iterations (int, optional): The maximum number of iterations to perform. Default is infinity.</span>
<span class="sd"> mask (ndarray, optional): A mask specifying which values to consider for filling NaNs. If None, all non-NaN values are used. Default is None.</span>
<span class="sd"> Returns:</span>
<span class="sd"> ndarray: The grid with NaN values filled using neighborhood averaging.</span>
<span class="sd"> Note:</span>
<span class="sd"> - This function requires the 'skimage' package, specifically the 'find_boundaries' function from 'skimage.segmentation'.</span>
<span class="sd"> Example:</span>
<span class="sd"> >>> grid = np.array([[1, 2, np.nan], [4, np.nan, 6], [np.nan, 8, 9]])</span>
<span class="sd"> >>> filled_grid = fill_nans(grid, neighborhood=2, max_iterations=10)</span>
<span class="sd"> >>> print(filled_grid)</span>
<span class="sd"> """</span>
<span class="kn">from</span><span class="w"> </span><span class="nn">skimage.segmentation</span><span class="w"> </span><span class="kn">import</span> <span class="n">find_boundaries</span>
<span class="n">ret</span> <span class="o">=</span> <span class="n">np</span><span class="o">.</span><span class="n">array</span><span class="p">(</span><span class="n">grid</span><span class="p">)</span>
<span class="k">if</span> <span class="n">mask</span> <span class="ow">is</span> <span class="kc">None</span><span class="p">:</span> <span class="n">mask</span> <span class="o">=</span> <span class="o">~</span><span class="n">np</span><span class="o">.</span><span class="n">isnan</span><span class="p">(</span><span class="n">grid</span><span class="p">)</span>
<span class="k">assert</span><span class="p">(</span><span class="n">np</span><span class="o">.</span><span class="n">all</span><span class="p">(</span><span class="n">np</span><span class="o">.</span><span class="n">isnan</span><span class="p">(</span><span class="n">mask</span><span class="p">))</span> <span class="o">==</span> <span class="mi">0</span><span class="p">)</span>
<span class="n">nans</span><span class="o">=</span> <span class="n">np</span><span class="o">.</span><span class="n">isnan</span><span class="p">(</span><span class="n">ret</span><span class="p">)</span>
<span class="n">i</span> <span class="o">=</span> <span class="mi">0</span>
<span class="k">while</span> <span class="n">np</span><span class="o">.</span><span class="n">sum</span><span class="p">(</span><span class="n">nans</span><span class="p">)</span> <span class="o">></span> <span class="mi">0</span><span class="p">:</span>
<span class="k">if</span> <span class="n">i</span> <span class="o">></span> <span class="n">max_iterations</span><span class="p">:</span> <span class="k">break</span>
<span class="n">i</span><span class="o">+=</span><span class="mi">1</span>
<span class="nb">print</span><span class="p">(</span><span class="s2">"nans"</span><span class="p">,</span><span class="n">np</span><span class="o">.</span><span class="n">sum</span><span class="p">(</span><span class="n">nans</span><span class="p">))</span>
<span class="n">boundary</span><span class="o">=</span><span class="n">find_boundaries</span><span class="p">(</span><span class="o">~</span><span class="n">nans</span><span class="p">,</span> <span class="n">mode</span><span class="o">=</span><span class="s1">'outer'</span><span class="p">)</span>
<span class="n">tmp</span> <span class="o">=</span> <span class="n">neighborhood_average</span><span class="p">(</span><span class="n">ret</span><span class="p">,</span> <span class="n">neighborhood</span><span class="p">)</span>
<span class="n">ret</span><span class="p">[</span><span class="n">boundary</span><span class="p">]</span> <span class="o">=</span> <span class="n">tmp</span><span class="p">[</span><span class="n">boundary</span><span class="p">]</span>
<span class="n">ret</span><span class="p">[</span><span class="n">mask</span><span class="p">]</span> <span class="o">=</span> <span class="n">grid</span><span class="p">[</span><span class="n">mask</span><span class="p">]</span>
<span class="n">nans</span> <span class="o">=</span> <span class="n">np</span><span class="o">.</span><span class="n">isnan</span><span class="p">(</span><span class="n">ret</span><span class="p">)</span>
<span class="k">return</span> <span class="n">ret</span></div>
<div class="viewcode-block" id="convolve_kernel_truncate">
<a class="viewcode-back" href="../../api/utilities/grid.html#arbdmodel.grid.convolve_kernel_truncate">[docs]</a>
<span class="k">def</span><span class="w"> </span><span class="nf">convolve_kernel_truncate</span><span class="p">(</span> <span class="n">array</span><span class="p">,</span> <span class="n">kernel</span> <span class="p">):</span>
<span class="w"> </span><span class="sd">"""</span>
<span class="sd"> Convolve an array with a kernel, truncating the output.</span>
<span class="sd"> Parameters:</span>
<span class="sd"> array (ndarray): The input array.</span>
<span class="sd"> kernel (ndarray): The kernel to convolve with the array.</span>
<span class="sd"> Returns:</span>
<span class="sd"> ndarray: The truncated convolution result.</span>
<span class="sd"> Raises:</span>
<span class="sd"> AssertionError: If the dimensions of the array and kernel do not match.</span>
<span class="sd"> AssertionError: If any dimension of the array is smaller than the corresponding dimension of the kernel.</span>
<span class="sd"> Note:</span>
<span class="sd"> - This function assumes that the kernel has odd dimensions along each dimension.</span>
<span class="sd"> - If any dimension of the kernel has an even number of elements, a warning message is printed to stderr indicating that the output may be shifted.</span>
<span class="sd"> Example:</span>
<span class="sd"> >>> array = np.array([[1, 2, 3], [4, 5, 6], [7, 8, 9]])</span>
<span class="sd"> >>> kernel = np.array([[0.5, 0.5], [0.5, 0.5]])</span>
<span class="sd"> >>> result = convolve_kernel_truncate(array, kernel)</span>
<span class="sd"> >>> print(result)</span>
<span class="sd"> """</span>
<span class="n">arrayShape</span> <span class="o">=</span> <span class="n">np</span><span class="o">.</span><span class="n">shape</span><span class="p">(</span> <span class="n">array</span> <span class="p">)</span>
<span class="n">kernelShape</span> <span class="o">=</span> <span class="n">np</span><span class="o">.</span><span class="n">shape</span><span class="p">(</span> <span class="n">kernel</span> <span class="p">)</span>
<span class="k">assert</span><span class="p">(</span> <span class="nb">len</span><span class="p">(</span><span class="n">arrayShape</span><span class="p">)</span> <span class="o">==</span> <span class="nb">len</span><span class="p">(</span><span class="n">kernelShape</span><span class="p">)</span> <span class="p">)</span>
<span class="k">for</span> <span class="n">an</span><span class="p">,</span><span class="n">kn</span> <span class="ow">in</span> <span class="nb">zip</span><span class="p">(</span><span class="n">arrayShape</span><span class="p">,</span><span class="n">kernelShape</span><span class="p">):</span> <span class="k">assert</span><span class="p">(</span><span class="n">an</span> <span class="o">></span> <span class="n">kn</span><span class="p">)</span>
<span class="n">dim</span> <span class="o">=</span> <span class="mi">0</span>
<span class="k">for</span> <span class="n">an</span> <span class="ow">in</span> <span class="n">kernelShape</span><span class="p">:</span>
<span class="n">dim</span> <span class="o">+=</span> <span class="mi">1</span>
<span class="k">if</span> <span class="n">an</span> <span class="o">%</span> <span class="mi">2</span> <span class="o">==</span> <span class="mi">0</span><span class="p">:</span>
<span class="nb">print</span><span class="p">(</span><span class="s2">"WARNING: kernel has even number of elements along dimension </span><span class="si">%d</span><span class="s2">; Output may be shifted."</span> <span class="o">%</span> <span class="n">dim</span><span class="p">,</span> <span class="n">file</span><span class="o">=</span><span class="n">sys</span><span class="o">.</span><span class="n">stderr</span><span class="p">)</span>
<span class="n">initShape</span> <span class="o">=</span> <span class="p">[</span><span class="n">a</span><span class="o">+</span><span class="n">b</span> <span class="k">for</span> <span class="n">a</span><span class="p">,</span><span class="n">b</span> <span class="ow">in</span> <span class="nb">zip</span><span class="p">(</span><span class="n">arrayShape</span><span class="p">,</span><span class="n">kernelShape</span><span class="p">)]</span>
<span class="n">convolution</span> <span class="o">=</span> <span class="n">np</span><span class="o">.</span><span class="n">zeros</span><span class="p">(</span><span class="n">initShape</span><span class="p">)</span>
<span class="n">count</span> <span class="o">=</span> <span class="n">np</span><span class="o">.</span><span class="n">zeros</span><span class="p">(</span><span class="n">initShape</span><span class="p">)</span>
<span class="k">for</span> <span class="n">idx</span> <span class="ow">in</span> <span class="nb">range</span><span class="p">(</span> <span class="n">np</span><span class="o">.</span><span class="n">prod</span><span class="p">(</span><span class="n">kernelShape</span><span class="p">)</span> <span class="p">):</span>
<span class="n">ijk</span> <span class="o">=</span> <span class="p">[</span> <span class="nb">int</span><span class="p">(</span><span class="n">idx</span><span class="o">/</span><span class="n">a</span><span class="p">)</span> <span class="o">%</span> <span class="n">b</span> <span class="k">for</span> <span class="n">a</span><span class="p">,</span><span class="n">b</span> <span class="ow">in</span>
<span class="nb">zip</span><span class="p">(</span> <span class="n">np</span><span class="o">.</span><span class="n">hstack</span><span class="p">(([</span><span class="mi">1</span><span class="p">],</span><span class="n">np</span><span class="o">.</span><span class="n">cumprod</span><span class="p">(</span><span class="n">kernelShape</span><span class="p">[:</span><span class="o">-</span><span class="mi">1</span><span class="p">]))),</span> <span class="n">kernelShape</span><span class="p">)</span> <span class="p">][::</span><span class="o">-</span><span class="mi">1</span><span class="p">]</span>
<span class="n">s</span> <span class="o">=</span> <span class="nb">tuple</span><span class="p">([</span> <span class="nb">slice</span><span class="p">(</span> <span class="p">(</span><span class="n">kn</span><span class="o">-</span><span class="n">i</span><span class="p">),</span> <span class="n">an</span><span class="o">+</span><span class="n">kn</span><span class="o">-</span><span class="n">i</span> <span class="p">)</span> <span class="k">for</span> <span class="n">i</span><span class="p">,</span><span class="n">kn</span><span class="p">,</span><span class="n">an</span> <span class="ow">in</span>
<span class="nb">zip</span><span class="p">(</span><span class="n">ijk</span><span class="p">,</span><span class="n">kernelShape</span><span class="p">,</span><span class="n">arrayShape</span><span class="p">)</span> <span class="p">])</span>
<span class="n">val</span> <span class="o">=</span> <span class="n">kernel</span><span class="o">.</span><span class="n">flatten</span><span class="p">()[</span><span class="n">idx</span><span class="p">]</span>
<span class="n">convolution</span><span class="p">[</span><span class="n">s</span><span class="p">]</span> <span class="o">+=</span> <span class="n">val</span><span class="o">*</span><span class="n">array</span>
<span class="n">count</span><span class="p">[</span><span class="n">s</span><span class="p">]</span> <span class="o">+=</span> <span class="mi">1</span>
<span class="n">ids</span> <span class="o">=</span> <span class="n">count</span> <span class="o">></span> <span class="mi">0</span>
<span class="n">convolution</span><span class="p">[</span><span class="n">ids</span><span class="p">]</span> <span class="o">=</span> <span class="p">(</span><span class="n">convolution</span><span class="p">[</span><span class="n">ids</span><span class="p">]</span><span class="o">/</span><span class="n">count</span><span class="p">[</span><span class="n">ids</span><span class="p">])</span> <span class="o">*</span> <span class="n">np</span><span class="o">.</span><span class="n">prod</span><span class="p">(</span><span class="n">kernelShape</span><span class="p">)</span>
<span class="n">s</span> <span class="o">=</span> <span class="nb">tuple</span><span class="p">([</span> <span class="nb">slice</span><span class="p">(</span> <span class="nb">int</span><span class="p">((</span><span class="n">kn</span><span class="o">-</span><span class="mi">1</span><span class="p">)</span><span class="o">/</span><span class="mi">2</span><span class="p">)</span><span class="o">+</span><span class="mi">1</span><span class="p">,</span> <span class="o">-</span><span class="nb">int</span><span class="p">((</span><span class="n">kn</span><span class="o">-</span><span class="mi">1</span><span class="p">)</span><span class="o">/</span><span class="mi">2</span><span class="p">)</span> <span class="p">)</span> <span class="k">for</span> <span class="n">kn</span> <span class="ow">in</span> <span class="n">kernelShape</span> <span class="p">])</span>
<span class="k">return</span> <span class="n">convolution</span><span class="p">[</span><span class="n">s</span><span class="p">]</span></div>
<div class="viewcode-block" id="isotropic_kernel">
<a class="viewcode-back" href="../../api/utilities/grid.html#arbdmodel.grid.isotropic_kernel">[docs]</a>
<span class="k">def</span><span class="w"> </span><span class="nf">isotropic_kernel</span><span class="p">(</span> <span class="n">function</span><span class="p">,</span> <span class="n">delta</span><span class="p">,</span> <span class="n">shape</span> <span class="o">=</span> <span class="kc">None</span><span class="p">,</span> <span class="n">cutoff</span> <span class="o">=</span> <span class="kc">None</span><span class="p">,</span> <span class="n">normalize</span> <span class="o">=</span> <span class="kc">False</span><span class="p">):</span>
<span class="w"> </span><span class="sd">"""</span>
<span class="sd"> Generate an isotropic kernel function.</span>
<span class="sd"> Parameters:</span>
<span class="sd"> function (callable): A callable representing the kernel function.</span>
<span class="sd"> It should accept a scalar or an array-like object and return a scalar or an array-like object of the same shape.</span>
<span class="sd"> delta (list or tuple): A list or tuple specifying the spacing between grid points in each dimension.</span>
<span class="sd"> shape (list or tuple, optional): A list or tuple specifying the shape of the kernel. If None, it is determined based on the cutoff parameter.</span>
<span class="sd"> cutoff (scalar, optional): A scalar specifying the maximum distance from the center to consider for the kernel. If provided, it determines the shape of the kernel based on the delta values.</span>
<span class="sd"> normalize (bool, optional): A boolean indicating whether to normalize the kernel. If True, the kernel values are divided by the sum of all kernel values, ensuring that the kernel sums to 1.</span>
<span class="sd"> Returns:</span>
<span class="sd"> ndarray: An ndarray representing the isotropic kernel.</span>
<span class="sd"> Raises:</span>
<span class="sd"> ValueError: If the cutoff parameter is None, as it is required to determine the shape of the kernel.</span>
<span class="sd"> Example:</span>
<span class="sd"> >>> import numpy as np</span>
<span class="sd"> >>> def gaussian(x):</span>
<span class="sd"> ... return np.exp(-0.5 * x**2)</span><span class="sd"> >>> delta = [0.1, 0.1, 0.1]</span>
<span class="sd"> >>> kernel = isotropic_kernel(gaussian, delta, cutoff=1.0, normalize=True)</span>
<span class="sd"> >>> print(kernel)</span>
<span class="sd"> """</span>
<span class="k">if</span> <span class="n">cutoff</span> <span class="ow">is</span> <span class="ow">not</span> <span class="kc">None</span><span class="p">:</span>
<span class="n">shape</span> <span class="o">=</span> <span class="p">[(</span><span class="n">cutoff</span><span class="o">*</span><span class="mi">2</span><span class="p">)</span><span class="o">//</span><span class="n">d</span> <span class="k">for</span> <span class="n">d</span> <span class="ow">in</span> <span class="n">delta</span><span class="p">]</span>
<span class="k">elif</span> <span class="n">cutoff</span> <span class="ow">is</span> <span class="kc">None</span><span class="p">:</span>
<span class="k">raise</span> <span class="ne">ValueError</span>
<span class="n">ndim</span> <span class="o">=</span> <span class="nb">len</span><span class="p">(</span><span class="n">shape</span><span class="p">)</span>
<span class="k">assert</span><span class="p">(</span> <span class="nb">len</span><span class="p">(</span><span class="n">delta</span><span class="p">)</span> <span class="o">==</span> <span class="n">ndim</span> <span class="p">)</span>
<span class="n">centers</span> <span class="o">=</span> <span class="p">[(</span><span class="n">np</span><span class="o">.</span><span class="n">arange</span><span class="p">(</span><span class="n">n</span><span class="p">)</span><span class="o">-</span><span class="p">(</span><span class="n">n</span><span class="o">-</span><span class="mi">1</span><span class="p">)</span><span class="o">*</span><span class="mf">0.5</span><span class="p">)</span><span class="o">*</span><span class="n">d</span> <span class="k">for</span> <span class="n">n</span><span class="p">,</span><span class="n">d</span> <span class="ow">in</span> <span class="nb">zip</span><span class="p">(</span><span class="n">shape</span><span class="p">,</span><span class="n">delta</span><span class="p">)]</span>
<span class="n">CENTERS</span> <span class="o">=</span> <span class="n">np</span><span class="o">.</span><span class="n">meshgrid</span><span class="p">(</span><span class="o">*</span><span class="n">centers</span><span class="p">,</span> <span class="n">indexing</span><span class="o">=</span><span class="s1">'ij'</span><span class="p">)</span>
<span class="n">R</span> <span class="o">=</span> <span class="n">np</span><span class="o">.</span><span class="n">sqrt</span><span class="p">(</span> <span class="nb">sum</span><span class="p">([</span><span class="n">C</span><span class="o">**</span><span class="mi">2</span> <span class="k">for</span> <span class="n">C</span> <span class="ow">in</span> <span class="n">CENTERS</span><span class="p">])</span> <span class="p">)</span>
<span class="n">kernel</span> <span class="o">=</span> <span class="n">function</span><span class="p">(</span><span class="n">R</span><span class="p">)</span>
<span class="k">if</span> <span class="n">normalize</span><span class="p">:</span>
<span class="n">kernel</span> <span class="o">=</span> <span class="n">kernel</span><span class="o">/</span><span class="n">np</span><span class="o">.</span><span class="n">sum</span><span class="p">(</span><span class="n">kernel</span><span class="p">)</span>
<span class="k">return</span> <span class="n">kernel</span></div>
<div class="viewcode-block" id="gaussian_kernel">
<a class="viewcode-back" href="../../api/utilities/grid.html#arbdmodel.grid.gaussian_kernel">[docs]</a>
<span class="k">def</span><span class="w"> </span><span class="nf">gaussian_kernel</span><span class="p">(</span><span class="n">voxels</span><span class="o">=</span><span class="mi">5</span><span class="p">,</span> <span class="n">sig</span><span class="o">=</span><span class="mf">1.</span><span class="p">,</span> <span class="n">ndim</span><span class="o">=</span><span class="mi">3</span><span class="p">):</span>
<span class="c1">## TODO: rewrite using isotropic_kernel</span>
<span class="w"> </span><span class="sd">"""</span>
<span class="sd"> creates gaussian kernel with side length `l` and a sigma of `sig`</span>
<span class="sd"> """</span>
<span class="n">gauss_list</span> <span class="o">=</span> <span class="p">[]</span>
<span class="k">try</span><span class="p">:</span>
<span class="n">sig</span><span class="p">[</span><span class="mi">0</span><span class="p">]</span>
<span class="k">except</span><span class="p">:</span>
<span class="n">sig</span> <span class="o">=</span> <span class="p">[</span><span class="n">sig</span><span class="p">]</span><span class="o">*</span><span class="n">ndim</span>
<span class="k">try</span><span class="p">:</span>
<span class="n">voxels</span><span class="p">[</span><span class="mi">0</span><span class="p">]</span>
<span class="k">except</span><span class="p">:</span>
<span class="n">voxels</span> <span class="o">=</span> <span class="p">[</span><span class="n">voxels</span><span class="p">]</span><span class="o">*</span><span class="n">ndim</span>
<span class="k">for</span> <span class="n">l</span><span class="p">,</span><span class="n">s</span> <span class="ow">in</span> <span class="nb">zip</span><span class="p">(</span><span class="n">voxels</span><span class="p">,</span><span class="n">sig</span><span class="p">):</span>
<span class="n">ax</span> <span class="o">=</span> <span class="n">np</span><span class="o">.</span><span class="n">linspace</span><span class="p">(</span><span class="o">-</span><span class="p">(</span><span class="n">l</span> <span class="o">-</span> <span class="mi">1</span><span class="p">)</span> <span class="o">/</span> <span class="mf">2.</span><span class="p">,</span> <span class="p">(</span><span class="n">l</span> <span class="o">-</span> <span class="mi">1</span><span class="p">)</span> <span class="o">/</span> <span class="mf">2.</span><span class="p">,</span> <span class="n">l</span><span class="p">)</span>
<span class="n">gauss_list</span><span class="o">.</span><span class="n">append</span><span class="p">(</span> <span class="n">np</span><span class="o">.</span><span class="n">exp</span><span class="p">(</span><span class="o">-</span><span class="mf">0.5</span> <span class="o">*</span> <span class="n">np</span><span class="o">.</span><span class="n">square</span><span class="p">(</span><span class="n">ax</span><span class="p">)</span> <span class="o">/</span> <span class="n">np</span><span class="o">.</span><span class="n">square</span><span class="p">(</span><span class="n">s</span><span class="p">))</span> <span class="p">)</span>
<span class="n">gauss_list2</span> <span class="o">=</span> <span class="p">[]</span>
<span class="n">kernel</span> <span class="o">=</span> <span class="n">gauss_list</span><span class="p">[</span><span class="mi">0</span><span class="p">]</span> <span class="c1"># np.empty((0))</span>
<span class="k">for</span> <span class="n">i</span><span class="p">,</span><span class="n">g</span> <span class="ow">in</span> <span class="nb">enumerate</span><span class="p">(</span><span class="n">gauss_list</span><span class="p">[</span><span class="mi">1</span><span class="p">:]):</span>
<span class="n">i</span><span class="o">=</span><span class="n">i</span><span class="o">+</span><span class="mi">1</span>
<span class="n">sl</span> <span class="o">=</span> <span class="nb">tuple</span><span class="p">(</span><span class="nb">slice</span><span class="p">(</span><span class="kc">None</span><span class="p">)</span> <span class="k">if</span> <span class="n">j</span><span class="o">==</span><span class="n">i</span> <span class="k">else</span> <span class="kc">None</span> <span class="k">for</span> <span class="n">j</span> <span class="ow">in</span> <span class="nb">range</span><span class="p">(</span><span class="n">i</span><span class="p">))</span>
<span class="n">kernel</span> <span class="o">=</span> <span class="n">kernel</span><span class="p">[</span><span class="o">...</span><span class="p">,</span><span class="kc">None</span><span class="p">]</span><span class="o">*</span><span class="n">g</span><span class="p">[</span><span class="n">sl</span><span class="p">]</span>
<span class="k">return</span> <span class="n">kernel</span> <span class="o">/</span> <span class="n">np</span><span class="o">.</span><span class="n">sum</span><span class="p">(</span><span class="n">kernel</span><span class="p">)</span></div>
<div class="viewcode-block" id="slab_potential_z">
<a class="viewcode-back" href="../../api/utilities/grid.html#arbdmodel.grid.slab_potential_z">[docs]</a>
<span class="k">def</span><span class="w"> </span><span class="nf">slab_potential_z</span><span class="p">(</span><span class="n">force_constant</span><span class="p">,</span> <span class="n">center</span><span class="p">,</span> <span class="n">dimensions</span><span class="p">,</span> <span class="n">resolution</span><span class="p">,</span> <span class="n">exponent</span><span class="o">=</span><span class="mi">2</span><span class="p">):</span>
<span class="w"> </span><span class="sd">"""</span>
<span class="sd"> Generate a potential energy field that depends only on the z-coordinate, creating a slab-like potential.</span>
<span class="sd"> Parameters</span>
<span class="sd"> ----------</span>
<span class="sd"> force_constant : float</span>
<span class="sd"> The strength of the potential (energy per unit distance^exponent).</span>
<span class="sd"> center : list or tuple</span>
<span class="sd"> The (x, y, z) coordinates of the center of the potential.</span>
<span class="sd"> dimensions : list or tuple</span>
<span class="sd"> The (x, y, z) dimensions of the grid in length units.</span>
<span class="sd"> resolution : float or list or tuple</span><span class="sd"> The grid spacing in length units. If a single value is provided, it's used for all dimensions.</span>
<span class="sd"> exponent : int, optional</span>
<span class="sd"> The exponent determining the shape of the potential. Default is 2 (harmonic potential).</span>
<span class="sd"> Returns</span>
<span class="sd"> -------</span>
<span class="sd"> numpy.ndarray</span>
<span class="sd"> A 3D array representing the potential energy field, where the energy depends only</span>
<span class="sd"> on the distance from the z-center raised to the specified exponent.</span>
<span class="sd"> Note:</span>
<span class="sd"> -----</span>
<span class="sd"> The potential U at each point is calculated as:</span>
<span class="sd"> U = (force_constant/exponent) * |z - center_z|^exponent</span>
<span class="sd"> """</span>
<span class="k">try</span><span class="p">:</span>
<span class="n">resolution</span><span class="p">[</span><span class="mi">0</span><span class="p">]</span>
<span class="k">except</span><span class="p">:</span>
<span class="n">resolution</span> <span class="o">=</span> <span class="mi">3</span><span class="o">*</span><span class="p">[</span><span class="n">resolution</span><span class="p">]</span>
<span class="n">n_voxels</span> <span class="o">=</span> <span class="p">[</span><span class="nb">int</span><span class="p">(</span><span class="n">np</span><span class="o">.</span><span class="n">ceil</span><span class="p">(</span><span class="n">w</span><span class="o">/</span><span class="n">r</span><span class="p">))</span> <span class="k">for</span> <span class="n">w</span><span class="p">,</span><span class="n">r</span> <span class="ow">in</span> <span class="nb">zip</span><span class="p">(</span><span class="n">dimensions</span><span class="p">,</span><span class="n">resolution</span><span class="p">)]</span>
<span class="n">x</span><span class="p">,</span><span class="n">y</span><span class="p">,</span><span class="n">z</span> <span class="o">=</span> <span class="p">[</span><span class="n">np</span><span class="o">.</span><span class="n">arange</span><span class="p">(</span><span class="n">n</span><span class="p">)</span><span class="o">*</span><span class="n">r</span><span class="o">-</span><span class="p">(</span><span class="n">c</span><span class="o">+</span><span class="p">(</span><span class="n">n</span><span class="o">/</span><span class="mf">2.0</span><span class="p">)</span><span class="o">*</span><span class="n">r</span><span class="p">)</span> <span class="k">for</span> <span class="n">n</span><span class="p">,</span><span class="n">r</span><span class="p">,</span><span class="n">c</span> <span class="ow">in</span> <span class="nb">zip</span><span class="p">(</span><span class="n">n_voxels</span><span class="p">,</span> <span class="n">resolution</span><span class="p">,</span> <span class="n">center</span><span class="p">)]</span>
<span class="n">u</span> <span class="o">=</span> <span class="p">(</span><span class="n">force_constant</span><span class="o">/</span><span class="n">exponent</span><span class="p">)</span> <span class="o">*</span> <span class="n">np</span><span class="o">.</span><span class="n">abs</span><span class="p">((</span><span class="n">z</span><span class="o">-</span><span class="n">center</span><span class="p">[</span><span class="mi">0</span><span class="p">])</span><span class="o">**</span><span class="n">exponent</span><span class="p">)</span> <span class="c1"># 1D array</span>
<span class="n">U</span> <span class="o">=</span> <span class="n">np</span><span class="o">.</span><span class="n">empty</span><span class="p">(</span> <span class="n">n_voxels</span> <span class="p">)</span>
<span class="n">U</span> <span class="o">=</span> <span class="n">u</span><span class="p">[</span><span class="kc">None</span><span class="p">,</span><span class="kc">None</span><span class="p">,:]</span>
<span class="k">return</span> <span class="n">U</span></div>
<div class="viewcode-block" id="constant_force">
<a class="viewcode-back" href="../../api/utilities/grid.html#arbdmodel.grid.constant_force">[docs]</a>
<span class="k">def</span><span class="w"> </span><span class="nf">constant_force</span><span class="p">(</span><span class="n">force</span><span class="p">,</span> <span class="n">dimensions</span><span class="p">,</span> <span class="n">resolution</span><span class="p">,</span> <span class="n">origin</span><span class="o">=</span><span class="kc">None</span><span class="p">):</span>
<span class="w"> </span><span class="sd">"""</span>
<span class="sd"> Generate a constant force field over a 3D grid.</span>
<span class="sd"> Parameters</span>
<span class="sd"> ----------</span>
<span class="sd"> force : float or list</span>
<span class="sd"> The force vector [fx, fy, fz] to apply. If a scalar is provided, </span>
<span class="sd"> it's interpreted as [0, 0, scalar].</span>
<span class="sd"> dimensions : list</span>
<span class="sd"> The physical dimensions [dx, dy, dz] of the grid in length units.</span>
<span class="sd"> resolution : float or list</span>
<span class="sd"> The resolution of each voxel [rx, ry, rz]. If a scalar is provided, </span>
<span class="sd"> it's used for all three dimensions.</span>
<span class="sd"> origin : list, optional</span>
<span class="sd"> The coordinates of the grid origin [ox, oy, oz]. If not provided, </span>
<span class="sd"> defaults to [-dx/2, -dy/2, -dz/2].</span>
<span class="sd"> Returns</span>
<span class="sd"> -------</span>
<span class="sd"> numpy.ndarray</span>
<span class="sd"> A 3D array representing the potential field U = fx*X + fy*Y + fz*Z, </span>
<span class="sd"> where X, Y, Z are the coordinate grids.</span>
<span class="sd"> """</span>
<span class="k">try</span><span class="p">:</span>
<span class="n">force</span><span class="p">[</span><span class="mi">0</span><span class="p">]</span>
<span class="k">except</span><span class="p">:</span>
<span class="n">force</span> <span class="o">=</span> <span class="p">[</span><span class="mi">0</span><span class="p">,</span><span class="mi">0</span><span class="p">,</span><span class="n">force</span><span class="p">]</span>
<span class="k">try</span><span class="p">:</span>
<span class="n">resolution</span><span class="p">[</span><span class="mi">0</span><span class="p">]</span>
<span class="k">except</span><span class="p">:</span>
<span class="n">resolution</span> <span class="o">=</span> <span class="mi">3</span><span class="o">*</span><span class="p">[</span><span class="n">resolution</span><span class="p">]</span>
<span class="k">if</span> <span class="n">origin</span> <span class="ow">is</span> <span class="kc">None</span><span class="p">:</span>
<span class="n">origin</span> <span class="o">=</span> <span class="p">[</span><span class="o">-</span><span class="mf">0.5</span> <span class="o">*</span> <span class="n">_x</span> <span class="k">for</span> <span class="n">_x</span> <span class="ow">in</span> <span class="n">dimensions</span><span class="p">]</span>
<span class="n">n_voxels</span> <span class="o">=</span> <span class="p">[</span><span class="nb">int</span><span class="p">(</span><span class="n">np</span><span class="o">.</span><span class="n">ceil</span><span class="p">(</span><span class="n">w</span><span class="o">/</span><span class="n">r</span><span class="p">))</span> <span class="k">for</span> <span class="n">w</span><span class="p">,</span><span class="n">r</span> <span class="ow">in</span> <span class="nb">zip</span><span class="p">(</span><span class="n">dimensions</span><span class="p">,</span><span class="n">resolution</span><span class="p">)]</span>
<span class="n">x</span><span class="p">,</span><span class="n">y</span><span class="p">,</span><span class="n">z</span> <span class="o">=</span> <span class="p">[(</span><span class="n">np</span><span class="o">.</span><span class="n">arange</span><span class="p">(</span><span class="n">n</span><span class="p">)</span><span class="o">+</span><span class="mf">0.5</span><span class="p">)</span><span class="o">*</span><span class="n">r</span><span class="o">+</span><span class="n">o</span> <span class="k">for</span> <span class="n">n</span><span class="p">,</span><span class="n">r</span><span class="p">,</span><span class="n">o</span> <span class="ow">in</span> <span class="nb">zip</span><span class="p">(</span><span class="n">n_voxels</span><span class="p">,</span> <span class="n">resolution</span><span class="p">,</span> <span class="n">origin</span><span class="p">)]</span>
<span class="n">X</span><span class="p">,</span><span class="n">Y</span><span class="p">,</span><span class="n">Z</span> <span class="o">=</span> <span class="n">np</span><span class="o">.</span><span class="n">meshgrid</span><span class="p">(</span><span class="n">x</span><span class="p">,</span><span class="n">y</span><span class="p">,</span><span class="n">z</span><span class="p">,</span><span class="n">indexing</span><span class="o">=</span><span class="s1">'ij'</span><span class="p">)</span>
<span class="n">U</span> <span class="o">=</span> <span class="n">force</span><span class="p">[</span><span class="mi">0</span><span class="p">]</span><span class="o">*</span><span class="n">X</span> <span class="o">+</span> <span class="n">force</span><span class="p">[</span><span class="mi">1</span><span class="p">]</span><span class="o">*</span><span class="n">Y</span> <span class="o">+</span> <span class="n">force</span><span class="p">[</span><span class="mi">2</span><span class="p">]</span> <span class="o">*</span> <span class="n">Z</span>
<span class="k">return</span> <span class="n">U</span></div>
<div class="viewcode-block" id="spherical_confinement">
<a class="viewcode-back" href="../../api/utilities/grid.html#arbdmodel.grid.spherical_confinement">[docs]</a>
<span class="k">def</span><span class="w"> </span><span class="nf">spherical_confinement</span><span class="p">(</span><span class="n">force_constant</span><span class="p">,</span> <span class="n">radius</span><span class="p">,</span> <span class="n">dimensions</span><span class="p">,</span> <span class="n">resolution</span><span class="p">,</span> <span class="n">center</span><span class="o">=</span><span class="p">(</span><span class="mi">0</span><span class="p">,</span><span class="mi">0</span><span class="p">,</span><span class="mi">0</span><span class="p">),</span> <span class="n">exponent</span><span class="o">=</span><span class="mi">2</span><span class="p">):</span>
<span class="w"> </span><span class="sd">"""</span>
<span class="sd"> Creates a spherical confinement potential field in 3D space.</span>
<span class="sd"> The potential is zero inside the sphere (R <= radius) and follows a power law </span>
<span class="sd"> (force_constant/exponent) * (R-radius)^exponent outside the sphere (R > radius).</span>
<span class="sd"> Parameters</span>
<span class="sd"> ----------</span>
<span class="sd"> force_constant : float</span>
<span class="sd"> The force constant that determines the strength of the confinement potential.</span>
<span class="sd"> radius : float</span>
<span class="sd"> Radius of the spherical confinement in the same units as dimensions and resolution.</span>
<span class="sd"> dimensions : list or tuple</span>
<span class="sd"> The dimensions of the 3D space [x_dim, y_dim, z_dim].</span>
<span class="sd"> resolution : float or list or tuple</span>
<span class="sd"> The spatial resolution of the grid. If a single value is provided, it will be used </span>
<span class="sd"> for all three dimensions. Otherwise, provide [x_res, y_res, z_res].</span>
<span class="sd"> center : tuple, optional</span>
<span class="sd"> The center coordinates of the sphere (x, y, z), default is (0, 0, 0).</span>
<span class="sd"> exponent : int, optional</span>
<span class="sd"> The exponent in the power law of the potential, default is 2.</span>
<span class="sd"> Returns</span>
<span class="sd"> -------</span>
<span class="sd"> U : numpy.ndarray</span>
<span class="sd"> 3D array containing the potential energy values at each grid point.</span>
<span class="sd"> Note</span>
<span class="sd"> -----</span>
<span class="sd"> The function seems to have an unused 'force' variable that might be part of incomplete code.</span>
<span class="sd"> """</span>
<span class="k">try</span><span class="p">:</span>
<span class="n">force</span><span class="p">[</span><span class="mi">0</span><span class="p">]</span>
<span class="k">except</span><span class="p">:</span>
<span class="n">force</span> <span class="o">=</span> <span class="p">[</span><span class="mi">0</span><span class="p">,</span><span class="mi">0</span><span class="p">,</span><span class="n">force</span><span class="p">]</span>
<span class="k">try</span><span class="p">:</span>
<span class="n">resolution</span><span class="p">[</span><span class="mi">0</span><span class="p">]</span>
<span class="k">except</span><span class="p">:</span>
<span class="n">resolution</span> <span class="o">=</span> <span class="mi">3</span><span class="o">*</span><span class="p">[</span><span class="n">resolution</span><span class="p">]</span>
<span class="n">n_voxels</span> <span class="o">=</span> <span class="p">[</span><span class="nb">int</span><span class="p">(</span><span class="n">np</span><span class="o">.</span><span class="n">ceil</span><span class="p">(</span><span class="n">w</span><span class="o">/</span><span class="n">r</span><span class="p">))</span> <span class="k">for</span> <span class="n">w</span><span class="p">,</span><span class="n">r</span> <span class="ow">in</span> <span class="nb">zip</span><span class="p">(</span><span class="n">dimensions</span><span class="p">,</span><span class="n">resolution</span><span class="p">)]</span>
<span class="n">x</span><span class="p">,</span><span class="n">y</span><span class="p">,</span><span class="n">z</span> <span class="o">=</span> <span class="p">[</span><span class="n">np</span><span class="o">.</span><span class="n">arange</span><span class="p">(</span><span class="n">n</span><span class="p">)</span><span class="o">*</span><span class="n">r</span><span class="o">-</span><span class="p">(</span><span class="n">c</span><span class="o">+</span><span class="p">(</span><span class="n">n</span><span class="o">/</span><span class="mf">2.0</span><span class="p">)</span><span class="o">*</span><span class="n">r</span><span class="p">)</span> <span class="k">for</span> <span class="n">n</span><span class="p">,</span><span class="n">r</span><span class="p">,</span><span class="n">c</span> <span class="ow">in</span> <span class="nb">zip</span><span class="p">(</span><span class="n">n_voxels</span><span class="p">,</span> <span class="n">resolution</span><span class="p">,</span> <span class="n">center</span><span class="p">)]</span>
<span class="n">X</span><span class="p">,</span><span class="n">Y</span><span class="p">,</span><span class="n">Z</span> <span class="o">=</span> <span class="n">np</span><span class="o">.</span><span class="n">meshgrid</span><span class="p">(</span><span class="n">x</span><span class="p">,</span><span class="n">y</span><span class="p">,</span><span class="n">z</span><span class="p">,</span><span class="n">indexing</span><span class="o">=</span><span class="s1">'ij'</span><span class="p">)</span>
<span class="n">R</span> <span class="o">=</span> <span class="n">np</span><span class="o">.</span><span class="n">sqrt</span><span class="p">(</span><span class="n">X</span><span class="o">**</span><span class="mi">2</span><span class="o">+</span><span class="n">Y</span><span class="o">**</span><span class="mi">2</span><span class="o">+</span><span class="n">Z</span><span class="o">**</span><span class="mi">2</span><span class="p">)</span>
<span class="n">U</span> <span class="o">=</span> <span class="n">np</span><span class="o">.</span><span class="n">empty</span><span class="p">(</span><span class="n">R</span><span class="o">.</span><span class="n">shape</span><span class="p">)</span>
<span class="n">sl</span> <span class="o">=</span> <span class="n">R</span> <span class="o">></span> <span class="n">radius</span>
<span class="n">U</span><span class="p">[</span><span class="n">sl</span><span class="p">]</span> <span class="o">=</span> <span class="p">(</span><span class="n">force_constant</span><span class="o">/</span><span class="n">exponent</span><span class="p">)</span> <span class="o">*</span> <span class="p">(</span><span class="n">R</span><span class="o">-</span><span class="n">radius</span><span class="p">)</span><span class="o">**</span><span class="n">exponent</span>
<span class="n">U</span><span class="p">[</span><span class="o">~</span><span class="n">sl</span><span class="p">]</span> <span class="o">=</span> <span class="mi">0</span>
<span class="k">return</span> <span class="n">U</span></div>
<div class="viewcode-block" id="write_confine_dx">
<a class="viewcode-back" href="../../api/utilities/grid.html#arbdmodel.grid.write_confine_dx">[docs]</a>
<span class="k">def</span><span class="w"> </span><span class="nf">write_confine_dx</span><span class="p">(</span><span class="n">radius</span><span class="o">=</span><span class="mi">100</span> <span class="p">):</span> <span class="c1">#Might merge with spherical confinement? </span>
<span class="w"> </span><span class="sd">"""</span>
<span class="sd"> Create and write a spherical confinement potential to a DX file.</span>
<span class="sd"> This function generates a spherical confinement potential with a harmonic potential</span>
<span class="sd"> outside the specified radius and a linear potential beyond radius + 25Å. The potential</span>
<span class="sd"> is written to a DX file named 'confine-{radius}.dx'.</span>
<span class="sd"> Parameters</span><span class="sd"> ----------</span>
<span class="sd"> radius : float, default=100</span>
<span class="sd"> Radius of the spherical confinement in Angstroms. The potential is zero within</span>
<span class="sd"> this radius and increases harmonically outside it.</span>
<span class="sd"> Returns</span>
<span class="sd"> -------</span>
<span class="sd"> None</span>
<span class="sd"> The function returns nothing but writes the potential to a DX file.</span>
<span class="sd"> If the file already exists, the function returns without doing anything.</span>
<span class="sd"> Note</span>
<span class="sd"> -----</span>
<span class="sd"> - The grid extends from -radius-50 to +radius+50 in all dimensions</span>
<span class="sd"> - Grid spacing is 2Å in all dimensions</span>
<span class="sd"> - The spring constant k is set to 1 kcal/mol/Ų</span>
<span class="sd"> - The potential transitions from harmonic to linear at radius + 25Å</span>
<span class="sd"> """</span>
<span class="n">outfile</span><span class="o">=</span><span class="sa">f</span><span class="s1">'confine-</span><span class="si">{</span><span class="n">radius</span><span class="si">}</span><span class="s1">.dx'</span>
<span class="k">if</span> <span class="n">Path</span><span class="p">(</span><span class="n">outfile</span><span class="p">)</span><span class="o">.</span><span class="n">exists</span><span class="p">():</span> <span class="k">return</span>
<span class="n">k</span> <span class="o">=</span> <span class="mi">1</span> <span class="c1"># Spring constant [kcal/mol/AA^2]</span>
<span class="c1"># radius=800</span>
<span class="n">x0</span> <span class="o">=</span> <span class="n">y0</span> <span class="o">=</span> <span class="o">-</span><span class="n">radius</span> <span class="o">-</span> <span class="mi">50</span>
<span class="n">x1</span> <span class="o">=</span> <span class="n">y1</span> <span class="o">=</span> <span class="o">-</span><span class="n">x0</span>
<span class="n">z0</span><span class="p">,</span><span class="n">z1</span> <span class="o">=</span> <span class="n">x0</span><span class="p">,</span><span class="n">x1</span>
<span class="n">dx</span> <span class="o">=</span> <span class="n">dy</span> <span class="o">=</span> <span class="n">dz</span> <span class="o">=</span> <span class="mi">2</span>
<span class="k">assert</span><span class="p">(</span> <span class="n">x1</span> <span class="o">></span> <span class="n">x0</span> <span class="p">)</span>
<span class="k">assert</span><span class="p">(</span> <span class="n">y1</span> <span class="o">></span> <span class="n">y0</span> <span class="p">)</span>
<span class="k">assert</span><span class="p">(</span> <span class="n">z1</span> <span class="o">></span> <span class="n">z0</span> <span class="p">)</span>
<span class="w"> </span><span class="sd">""" Create grid axes """</span>
<span class="n">x</span><span class="p">,</span><span class="n">y</span><span class="p">,</span><span class="n">z</span> <span class="o">=</span> <span class="p">[</span><span class="n">np</span><span class="o">.</span><span class="n">arange</span><span class="p">(</span> <span class="n">a</span><span class="o">-</span><span class="n">res</span><span class="o">/</span><span class="mi">2</span><span class="p">,</span> <span class="n">b</span><span class="o">+</span><span class="n">res</span><span class="o">/</span><span class="mi">2</span><span class="p">,</span> <span class="n">res</span> <span class="p">)</span>
<span class="k">for</span> <span class="n">a</span><span class="p">,</span><span class="n">b</span><span class="p">,</span><span class="n">res</span> <span class="ow">in</span> <span class="nb">zip</span><span class="p">((</span><span class="n">x0</span><span class="p">,</span><span class="n">y0</span><span class="p">,</span><span class="n">z0</span><span class="p">),(</span><span class="n">x1</span><span class="p">,</span><span class="n">y1</span><span class="p">,</span><span class="n">z1</span><span class="p">),(</span><span class="n">dx</span><span class="p">,</span><span class="n">dy</span><span class="p">,</span><span class="n">dz</span><span class="p">))]</span>
<span class="c1"># x = np.arange( -100, 100, dx ) # alternatively, be explicit</span>
<span class="c1"># assert( x[0] == -x[-1] )</span>
<span class="n">X</span><span class="p">,</span><span class="n">Y</span><span class="p">,</span><span class="n">Z</span> <span class="o">=</span> <span class="n">np</span><span class="o">.</span><span class="n">meshgrid</span><span class="p">(</span><span class="n">x</span><span class="p">,</span><span class="n">y</span><span class="p">,</span><span class="n">z</span><span class="p">,</span><span class="n">indexing</span><span class="o">=</span><span class="s1">'ij'</span><span class="p">)</span> <span class="c1"># create meshgrid for making potential</span>
<span class="n">R</span> <span class="o">=</span> <span class="n">np</span><span class="o">.</span><span class="n">sqrt</span><span class="p">(</span><span class="n">X</span><span class="o">**</span><span class="mi">2</span> <span class="o">+</span> <span class="n">Y</span><span class="o">**</span><span class="mi">2</span> <span class="o">+</span> <span class="n">Z</span><span class="o">**</span><span class="mi">2</span><span class="p">)</span>
<span class="w"> </span><span class="sd">""" Create the potential, adding 0.5 k deltaX**2 for each half plane """</span>
<span class="n">pot</span> <span class="o">=</span> <span class="n">np</span><span class="o">.</span><span class="n">zeros</span><span class="p">(</span> <span class="n">X</span><span class="o">.</span><span class="n">shape</span> <span class="p">)</span>
<span class="n">ids</span> <span class="o">=</span> <span class="n">R</span> <span class="o">></span> <span class="n">radius</span>
<span class="n">pot</span><span class="p">[</span><span class="n">ids</span><span class="p">]</span> <span class="o">=</span> <span class="mf">0.5</span><span class="o">*</span><span class="n">k</span><span class="o">*</span><span class="p">(</span><span class="n">R</span><span class="p">[</span><span class="n">ids</span><span class="p">]</span><span class="o">-</span><span class="n">radius</span><span class="p">)</span><span class="o">**</span><span class="mi">2</span>
<span class="n">ids</span> <span class="o">=</span> <span class="n">R</span> <span class="o">></span> <span class="n">radius</span> <span class="o">+</span> <span class="mi">25</span>
<span class="n">pot</span><span class="p">[</span><span class="n">ids</span><span class="p">]</span> <span class="o">=</span> <span class="mf">0.5</span><span class="o">*</span><span class="n">k</span><span class="o">*</span><span class="mi">25</span><span class="o">**</span><span class="mi">2</span> <span class="o">+</span> <span class="mf">0.5</span><span class="o">*</span><span class="n">k</span><span class="o">*</span><span class="mi">25</span><span class="o">**</span><span class="mi">2</span><span class="o">*</span><span class="p">(</span><span class="n">R</span><span class="p">[</span><span class="n">ids</span><span class="p">]</span><span class="o">-</span><span class="n">radius</span><span class="o">-</span><span class="mi">25</span><span class="p">)</span> <span class="c1"># switch to linear potential</span>
<span class="w"> </span><span class="sd">""" Write the dx file """</span>
<span class="n">writeDx</span><span class="p">(</span><span class="n">outfile</span><span class="p">,</span> <span class="n">pot</span><span class="p">,</span>
<span class="n">delta</span><span class="o">=</span><span class="p">(</span><span class="n">dx</span><span class="p">,</span><span class="n">dy</span><span class="p">,</span><span class="n">dz</span><span class="p">),</span>
<span class="n">origin</span><span class="o">=</span><span class="p">(</span><span class="n">x</span><span class="p">[</span><span class="mi">0</span><span class="p">],</span><span class="n">y</span><span class="p">[</span><span class="mi">0</span><span class="p">],</span><span class="n">z</span><span class="p">[</span><span class="mi">0</span><span class="p">]))</span></div>
<span class="c1">## Utility functions</span>
<div class="viewcode-block" id="create_bounding_grid">
<a class="viewcode-back" href="../../api/utilities/grid.html#arbdmodel.grid.create_bounding_grid">[docs]</a>
<span class="k">def</span><span class="w"> </span><span class="nf">create_bounding_grid</span><span class="p">(</span> <span class="o">*</span><span class="n">grids</span> <span class="p">):</span>
<span class="w"> </span><span class="sd">""" Construct a grid bigger than all the (GridDataFormats) grids provided in the arguments; note that the inputs must be orthonormal and have exactly overlapping voxels """</span>
<span class="k">def</span><span class="w"> </span><span class="nf">_create_bounding_grid</span><span class="p">(</span> <span class="n">grid1</span><span class="p">,</span> <span class="n">grid2</span> <span class="p">):</span>
<span class="k">assert</span><span class="p">(</span> <span class="nb">len</span><span class="p">(</span><span class="n">grid1</span><span class="o">.</span><span class="n">grid</span><span class="o">.</span><span class="n">shape</span><span class="p">)</span> <span class="o">==</span> <span class="mi">3</span> <span class="p">)</span> <span class="k">assert</span><span class="p">(</span> <span class="nb">len</span><span class="p">(</span><span class="n">grid2</span><span class="o">.</span><span class="n">grid</span><span class="o">.</span><span class="n">shape</span><span class="p">)</span> <span class="o">==</span> <span class="mi">3</span> <span class="p">)</span>
<span class="k">assert</span><span class="p">(</span><span class="nb">len</span><span class="p">(</span><span class="n">grid1</span><span class="o">.</span><span class="n">delta</span><span class="p">)</span> <span class="o">==</span> <span class="mi">3</span><span class="p">)</span>
<span class="k">assert</span><span class="p">(</span> <span class="n">np</span><span class="o">.</span><span class="n">all</span><span class="p">(</span><span class="n">grid1</span><span class="o">.</span><span class="n">delta</span> <span class="o">==</span> <span class="n">grid2</span><span class="o">.</span><span class="n">delta</span><span class="p">)</span> <span class="p">)</span>
<span class="n">min_origin</span> <span class="o">=</span> <span class="n">np</span><span class="o">.</span><span class="n">minimum</span><span class="p">(</span><span class="o">*</span><span class="p">[</span><span class="n">g</span><span class="o">.</span><span class="n">origin</span> <span class="k">for</span> <span class="n">g</span> <span class="ow">in</span> <span class="p">(</span><span class="n">grid1</span><span class="p">,</span> <span class="n">grid2</span><span class="p">)])</span>
<span class="n">max_shape</span> <span class="o">=</span> <span class="n">np</span><span class="o">.</span><span class="n">maximum</span><span class="p">(</span><span class="o">*</span><span class="p">[(</span><span class="n">g</span><span class="o">.</span><span class="n">origin</span><span class="o">-</span><span class="n">min_origin</span><span class="p">)</span><span class="o">//</span><span class="n">g</span><span class="o">.</span><span class="n">delta</span> <span class="o">+</span> <span class="n">np</span><span class="o">.</span><span class="n">array</span><span class="p">(</span><span class="n">g</span><span class="o">.</span><span class="n">grid</span><span class="o">.</span><span class="n">shape</span><span class="p">)</span> <span class="k">for</span> <span class="n">g</span> <span class="ow">in</span> <span class="p">(</span><span class="n">grid1</span><span class="p">,</span> <span class="n">grid2</span><span class="p">)])</span><span class="o">.</span><span class="n">astype</span><span class="p">(</span><span class="nb">int</span><span class="p">)</span>
<span class="k">assert</span><span class="p">(</span> <span class="n">np</span><span class="o">.</span><span class="n">all</span><span class="p">(</span> <span class="n">max_shape</span> <span class="o">></span> <span class="mi">0</span><span class="p">)</span> <span class="p">)</span>
<span class="n">new_grid_data</span> <span class="o">=</span> <span class="n">np</span><span class="o">.</span><span class="n">empty</span><span class="p">(</span><span class="n">max_shape</span><span class="p">)</span>
<span class="n">new_grid</span> <span class="o">=</span> <span class="n">Grid</span><span class="p">(</span> <span class="n">new_grid_data</span><span class="p">,</span> <span class="n">origin</span> <span class="o">=</span> <span class="n">min_origin</span><span class="p">,</span> <span class="n">delta</span> <span class="o">=</span> <span class="n">grid1</span><span class="o">.</span><span class="n">delta</span> <span class="p">)</span>
<span class="k">return</span> <span class="n">new_grid</span>
<span class="n">g1</span> <span class="o">=</span> <span class="n">grids</span><span class="p">[</span><span class="mi">0</span><span class="p">]</span>
<span class="k">if</span> <span class="nb">len</span><span class="p">(</span> <span class="n">grids</span> <span class="p">)</span> <span class="o">==</span> <span class="mi">1</span><span class="p">:</span>
<span class="n">g1</span> <span class="o">=</span> <span class="n">_create_bounding_grid</span><span class="p">(</span> <span class="n">g1</span><span class="p">,</span> <span class="n">g1</span> <span class="p">)</span>
<span class="k">else</span><span class="p">:</span>
<span class="k">for</span> <span class="n">g2</span> <span class="ow">in</span> <span class="n">grids</span><span class="p">[</span><span class="mi">1</span><span class="p">:]:</span>
<span class="n">g1</span> <span class="o">=</span> <span class="n">_create_bounding_grid</span><span class="p">(</span> <span class="n">g1</span><span class="p">,</span> <span class="n">g2</span> <span class="p">)</span>
<span class="n">g1</span><span class="o">.</span><span class="n">grid</span><span class="p">[</span><span class="o">...</span><span class="p">,:]</span> <span class="o">=</span> <span class="mi">0</span>
<span class="k">return</span> <span class="n">g1</span></div>
<div class="viewcode-block" id="get_slice_enclosing_smaller_grid">
<a class="viewcode-back" href="../../api/utilities/grid.html#arbdmodel.grid.get_slice_enclosing_smaller_grid">[docs]</a>
<span class="k">def</span><span class="w"> </span><span class="nf">get_slice_enclosing_smaller_grid</span><span class="p">(</span> <span class="n">grid</span><span class="p">,</span> <span class="n">smaller_grid</span> <span class="p">):</span>
<span class="w"> </span><span class="sd">"""</span>
<span class="sd"> Calculates the slices necessary to extract a portion of a larger grid that encloses a smaller grid.</span>
<span class="sd"> This function computes the slice indices needed to extract a subgrid from a larger grid</span>
<span class="sd"> that corresponds to the region covered by a smaller grid. The function assumes that both</span>
<span class="sd"> grids are orthonormal and have the same grid spacing (delta).</span>
<span class="sd"> Parameters</span>
<span class="sd"> ----------</span>
<span class="sd"> grid : Grid</span>
<span class="sd"> The larger grid object containing a 3D array and grid metadata.</span>
<span class="sd"> Must have attributes: grid (3D numpy array), origin (3D coordinates), and delta (grid spacing).</span>
<span class="sd"> smaller_grid : Grid</span>
<span class="sd"> The smaller grid object that is fully contained within the larger grid.</span>
<span class="sd"> Must have the same attributes as grid.</span>
<span class="sd"> Returns</span>
<span class="sd"> -------</span>
<span class="sd"> list of slice</span>
<span class="sd"> A list of three slice objects, one for each dimension, that can be used to extract</span>
<span class="sd"> the portion of the larger grid that corresponds to the smaller grid.</span>
<span class="sd"> Raises</span>
<span class="sd"> ------</span>
<span class="sd"> AssertionError</span>
<span class="sd"> If grids are not 3D, if grid spacings are not identical, if smaller_grid is not</span>
<span class="sd"> contained within grid, if grids are not aligned on grid points, or if smaller_grid</span>
<span class="sd"> extends beyond the boundaries of grid.</span>
<span class="sd"> Note</span>
<span class="sd"> -----</span>
<span class="sd"> The function assumes that both grids are aligned such that the origin of smaller_grid</span>
<span class="sd"> falls exactly on a grid point of the larger grid.</span>
<span class="sd"> """</span>
<span class="c1">## Check grids are orthonormal</span>
<span class="k">assert</span><span class="p">(</span> <span class="nb">len</span><span class="p">(</span><span class="n">grid</span><span class="o">.</span><span class="n">grid</span><span class="o">.</span><span class="n">shape</span><span class="p">)</span> <span class="o">==</span> <span class="mi">3</span> <span class="p">)</span>
<span class="k">assert</span><span class="p">(</span> <span class="nb">len</span><span class="p">(</span><span class="n">smaller_grid</span><span class="o">.</span><span class="n">grid</span><span class="o">.</span><span class="n">shape</span><span class="p">)</span> <span class="o">==</span> <span class="mi">3</span> <span class="p">)</span>
<span class="k">assert</span><span class="p">(</span><span class="nb">len</span><span class="p">(</span><span class="n">grid</span><span class="o">.</span><span class="n">delta</span><span class="p">)</span> <span class="o">==</span> <span class="mi">3</span><span class="p">)</span>
<span class="c1">## Make sure grid spacing is the same</span>
<span class="k">assert</span><span class="p">(</span> <span class="n">np</span><span class="o">.</span><span class="n">all</span><span class="p">(</span><span class="n">grid</span><span class="o">.</span><span class="n">delta</span> <span class="o">==</span> <span class="n">smaller_grid</span><span class="o">.</span><span class="n">delta</span><span class="p">)</span> <span class="p">)</span>
<span class="n">delta</span> <span class="o">=</span> <span class="n">smaller_grid</span><span class="o">.</span><span class="n">delta</span>
<span class="c1">## Find the offset between the grids and make sure smaller_grid is inside grid</span>
<span class="n">offset</span> <span class="o">=</span> <span class="p">(</span><span class="n">smaller_grid</span><span class="o">.</span><span class="n">origin</span> <span class="o">-</span> <span class="n">grid</span><span class="o">.</span><span class="n">origin</span><span class="p">)</span> <span class="k">assert</span><span class="p">(</span> <span class="n">np</span><span class="o">.</span><span class="n">all</span><span class="p">(</span> <span class="n">offset</span> <span class="o">></span> <span class="mi">0</span> <span class="p">)</span> <span class="p">)</span>
<span class="c1">## Check that grids are overlapping exactly and find index offset</span>
<span class="n">slices</span> <span class="o">=</span> <span class="p">[]</span>
<span class="k">for</span> <span class="n">x</span><span class="p">,</span><span class="n">d</span><span class="p">,</span><span class="n">n</span><span class="p">,</span><span class="n">ln</span> <span class="ow">in</span> <span class="nb">zip</span><span class="p">(</span><span class="n">offset</span><span class="p">,</span> <span class="n">delta</span><span class="p">,</span> <span class="n">smaller_grid</span><span class="o">.</span><span class="n">grid</span><span class="o">.</span><span class="n">shape</span><span class="p">,</span> <span class="n">grid</span><span class="o">.</span><span class="n">grid</span><span class="o">.</span><span class="n">shape</span><span class="p">):</span>
<span class="k">assert</span><span class="p">(</span> <span class="n">x</span><span class="o">/</span><span class="n">d</span> <span class="o">==</span> <span class="n">x</span><span class="o">//</span><span class="n">d</span> <span class="p">)</span>
<span class="n">i0</span> <span class="o">=</span> <span class="nb">int</span><span class="p">(</span><span class="nb">round</span><span class="p">(</span><span class="n">x</span><span class="o">/</span><span class="n">d</span><span class="p">))</span>
<span class="k">assert</span><span class="p">(</span><span class="n">i0</span><span class="o">+</span><span class="n">n</span> <span class="o"><=</span> <span class="n">ln</span><span class="p">)</span> <span class="c1"># make sure that smaller_grid isn't too big for grid</span>
<span class="n">slices</span><span class="o">.</span><span class="n">append</span><span class="p">(</span><span class="nb">slice</span><span class="p">(</span><span class="n">i0</span><span class="p">,</span><span class="n">i0</span><span class="o">+</span><span class="n">n</span><span class="p">))</span>
<span class="k">return</span> <span class="n">slices</span></div>
<span class="k">if</span> <span class="vm">__name__</span> <span class="o">==</span> <span class="s1">'__main__'</span><span class="p">:</span>
<span class="n">unittest</span><span class="o">.</span><span class="n">main</span><span class="p">()</span>
</pre></div>
</article>
<footer class="prev-next-footer d-print-none">
<div class="prev-next-area">
</div>
</footer>
</div>
</div>
<footer class="bd-footer-content">
<div class="bd-footer-content__inner container">
<div class="footer-item">
<p class="component-author">
By ARBD Model Team
</p>
</div>
<div class="footer-item">
<p class="copyright">
© Copyright 2023.
<br/>
</p>
</div>
<div class="footer-item">
</div>
<div class="footer-item">
</div>
</div> </footer>
</main>
</div>
</div>
<!-- Scripts loaded after <body> so the DOM is not blocked -->
<script src="../../_static/scripts/bootstrap.js?digest=dfe6caa3a7d634c4db9b"></script>
<script src="../../_static/scripts/pydata-sphinx-theme.js?digest=dfe6caa3a7d634c4db9b"></script>
<footer class="bd-footer">
</footer>
</body>
</html>