From 73de5cd8b6c155d48ccd45e52449278e6c00657e Mon Sep 17 00:00:00 2001 From: Colm Talbot Date: Mon, 14 Nov 2022 15:18:45 -0500 Subject: [PATCH 1/4] MAINT: refactor package and add precommits --- .pre-commit-config.yaml | 40 ++++ bin/generate_gamma_data.py | 7 +- docs/conf.py | 74 +++--- examples/Comparison.ipynb | 10 +- examples/GWMemory.ipynb | 369 +++++++++++++++++------------- gwmemory/__init__.py | 2 +- gwmemory/angles.py | 108 ++++++--- gwmemory/gwmemory.py | 47 ++-- gwmemory/harmonics.py | 21 +- gwmemory/qnms.py | 13 +- gwmemory/utils.py | 19 +- gwmemory/waveforms/__init__.py | 2 +- gwmemory/waveforms/approximant.py | 72 +++--- gwmemory/waveforms/base.py | 34 +-- gwmemory/waveforms/mwm.py | 55 +++-- gwmemory/waveforms/nr.py | 29 ++- gwmemory/waveforms/surrogate.py | 45 ++-- pyproject.toml | 2 +- test/angles_test.py | 16 +- test/harmonics_test.py | 4 +- test/waveform_test.py | 16 +- 21 files changed, 596 insertions(+), 389 deletions(-) create mode 100644 .pre-commit-config.yaml diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml new file mode 100644 index 0000000..c1d3dbf --- /dev/null +++ b/.pre-commit-config.yaml @@ -0,0 +1,40 @@ +repos: +- repo: https://github.com/ambv/black + rev: 22.10.0 + hooks: + - id: black + language_version: python3 +- repo: https://github.com/ambv/black + rev: 22.10.0 + hooks: + - id: black-jupyter + language_version: python3 +- repo: https://github.com/pre-commit/pre-commit-hooks + rev: v4.3.0 + hooks: + - id: check-merge-conflict # prevent committing files with merge conflicts +- repo: https://github.com/codespell-project/codespell + rev: v2.2.2 + hooks: + - id: codespell # Spellchecker + args: [-L, "nd,hist", "--skip", "*.ipynb"] +- repo: https://github.com/pre-commit/mirrors-isort + rev: v5.10.1 + hooks: + - id: isort # sort imports alphabetically and separates import into sections + args: [-w=88, -m=3, --tc, -sp=setup.cfg ] +- repo: local + hooks: + - id: flynt + name: flynt + entry: flynt + args: [--fail-on-change] + types: [python] + language: python + additional_dependencies: + - flynt +- repo: https://github.com/ColmTalbot/nbstrip + rev: v0.1.1 + hooks: + - id: nbstrip + args: [-i, -v] diff --git a/bin/generate_gamma_data.py b/bin/generate_gamma_data.py index 5ca9996..758cef4 100644 --- a/bin/generate_gamma_data.py +++ b/bin/generate_gamma_data.py @@ -30,10 +30,9 @@ label_1 = f"{ell1}{m1}" for ell2, m2 in lms: label_2 = f"{ell2}{m2}" - gammalmlm = np.array([ - angles.analytic_gamma((ell1, m1), (ell2, m2), ell) - for ell in range(2, 21) - ]) + gammalmlm = np.array( + [angles.analytic_gamma((ell1, m1), (ell2, m2), ell) for ell in range(2, 21)] + ) delta_m = m1 - m2 new[str(delta_m)][label_1 + label_2] = np.real(gammalmlm) for key in new: diff --git a/docs/conf.py b/docs/conf.py index 6024285..ccc99ee 100644 --- a/docs/conf.py +++ b/docs/conf.py @@ -18,8 +18,10 @@ # import os import sys + import gwmemory -sys.path.insert(0, os.path.abspath('../')) + +sys.path.insert(0, os.path.abspath("../")) # -- General configuration ------------------------------------------------ @@ -30,25 +32,30 @@ # Add any Sphinx extension module names here, as strings. They can be # extensions coming with Sphinx (named 'sphinx.ext.*') or your custom # ones. -extensions = ['sphinx.ext.autodoc', 'sphinx.ext.mathjax', 'numpydoc', - 'nbsphinx', 'sphinx.ext.autosummary'] +extensions = [ + "sphinx.ext.autodoc", + "sphinx.ext.mathjax", + "numpydoc", + "nbsphinx", + "sphinx.ext.autosummary", +] # Add any paths that contain templates here, relative to this directory. -templates_path = ['_templates'] +templates_path = ["_templates"] # The suffix(es) of source filenames. # You can specify multiple suffix as a list of string: # -source_suffix = ['.rst', '.md', '.txt'] +source_suffix = [".rst", ".md", ".txt"] # source_suffix = '.txt' # The master toctree document. -master_doc = 'index' +master_doc = "index" # General information about the project. -project = u'GWMemory' -copyright = u'2018, Colm Talbot' -author = u'Colm Talbot' +project = "GWMemory" +copyright = "2018, Colm Talbot" +author = "Colm Talbot" # The version info for the project you're documenting, acts as replacement for @@ -56,8 +63,8 @@ # built documents. # # The short X.Y version. -fullversion = gwmemory.__version__.split(':')[0] -version = '.'.join(fullversion.split('.')[:2]) +fullversion = gwmemory.__version__.split(":")[0] +version = ".".join(fullversion.split(".")[:2]) # The full version, including alpha/beta/rc tags. release = fullversion @@ -71,10 +78,10 @@ # List of patterns, relative to source directory, that match files and # directories to ignore when looking for source files. # This patterns also effect to html_static_path and html_extra_path -exclude_patterns = ['_build', 'Thumbs.db', '.DS_Store', 'requirements.txt'] +exclude_patterns = ["_build", "Thumbs.db", ".DS_Store", "requirements.txt"] # The name of the Pygments (syntax highlighting) style to use. -pygments_style = 'sphinx' +pygments_style = "sphinx" # If true, `todo` and `todoList` produce output, else they produce nothing. todo_include_todos = False @@ -85,7 +92,7 @@ # The theme to use for HTML and HTML Help pages. See the documentation for # a list of builtin themes. # -html_theme = 'sphinx_rtd_theme' +html_theme = "sphinx_rtd_theme" # Theme options are theme-specific and customize the look and feel of a theme # further. For a list of options available for each theme, see the @@ -96,7 +103,7 @@ # Add any paths that contain custom static files (such as style sheets) here, # relative to this directory. They are copied after the builtin static files, # so a file named "default.css" will overwrite the builtin "default.css". -html_static_path = ['_static'] +html_static_path = ["_static"] # Custom sidebar templates, must be a dictionary that maps document names # to template names. @@ -104,12 +111,12 @@ # This is required for the alabaster theme # refs: http://alabaster.readthedocs.io/en/latest/installation.html#sidebars html_sidebars = { - '**': [ - 'about.html', - 'navigation.html', - 'relations.html', # needs 'show_related': True theme option to display - 'searchbox.html', - 'donate.html', + "**": [ + "about.html", + "navigation.html", + "relations.html", # needs 'show_related': True theme option to display + "searchbox.html", + "donate.html", ] } @@ -117,7 +124,7 @@ # -- Options for HTMLHelp output ------------------------------------------ # Output file base name for HTML help builder. -htmlhelp_basename = 'gwmemorydoc' +htmlhelp_basename = "gwmemorydoc" # -- Options for LaTeX output --------------------------------------------- @@ -126,15 +133,12 @@ # The paper size ('letterpaper' or 'a4paper'). # # 'papersize': 'letterpaper', - # The font size ('10pt', '11pt' or '12pt'). # # 'pointsize': '10pt', - # Additional stuff for the LaTeX preamble. # # 'preamble': '', - # Latex figure (float) alignment # # 'figure_align': 'htbp', @@ -144,8 +148,7 @@ # (source start file, target name, title, # author, documentclass [howto, manual, or own class]). latex_documents = [ - (master_doc, 'gwmemory.tex', u'GWMemory Documentation', - u'Colm Talbot', 'manual'), + (master_doc, "gwmemory.tex", "GWMemory Documentation", "Colm Talbot", "manual"), ] @@ -153,10 +156,7 @@ # One entry per manual page. List of tuples # (source start file, name, description, authors, manual section). -man_pages = [ - (master_doc, 'GWMemory', u'GWMemory Documentation', - [author], 1) -] +man_pages = [(master_doc, "GWMemory", "GWMemory Documentation", [author], 1)] # -- Options for Texinfo output ------------------------------------------- @@ -165,9 +165,15 @@ # (source start file, target name, title, author, # dir menu entry, description, category) texinfo_documents = [ - (master_doc, 'GWMemory', u'GWMemory Documentation', - author, 'Colm Talbot', 'Arbitrary gravitational-wave memory waveforms.', - 'Miscellaneous'), + ( + master_doc, + "GWMemory", + "GWMemory Documentation", + author, + "Colm Talbot", + "Arbitrary gravitational-wave memory waveforms.", + "Miscellaneous", + ), ] numpydoc_show_class_members = False diff --git a/examples/Comparison.ipynb b/examples/Comparison.ipynb index e96ca16..8776c6c 100644 --- a/examples/Comparison.ipynb +++ b/examples/Comparison.ipynb @@ -105,13 +105,19 @@ "for ii, mode in enumerate(modes):\n", " gwmem = h_mem[mode]\n", " sxsmem = h_mem_sxs[mode]\n", - " overlap = np.vdot(gwmem, sxsmem) / np.vdot(gwmem, gwmem) ** 0.5 / np.vdot(sxsmem, sxsmem) ** 0.5\n", + " overlap = (\n", + " np.vdot(gwmem, sxsmem)\n", + " / np.vdot(gwmem, gwmem) ** 0.5\n", + " / np.vdot(sxsmem, sxsmem) ** 0.5\n", + " )\n", "\n", " ax = axes[ii // 3, ii % 3]\n", " ax.plot(times, gwmem.real, label=\"GWMemory\")\n", " ax.plot(times_sxs, sxsmem.real, linestyle=\":\", label=\"SXS\")\n", " ax.annotate(f\"Mode: ({mode[0]}, {mode[1]})\", (0.03, 0.75), xycoords=\"axes fraction\")\n", - " ax.annotate(f\"Mismatch: {1 - overlap.real:.1e}\", (0.35, 0.75), xycoords=\"axes fraction\")\n", + " ax.annotate(\n", + " f\"Mismatch: {1 - overlap.real:.1e}\", (0.35, 0.75), xycoords=\"axes fraction\"\n", + " )\n", "axes[0, 0].legend(loc=\"lower left\")\n", "axes[-1, 0].set_xlabel(\"$t [M]$\")\n", "axes[-1, 1].set_xlabel(\"$t [M]$\")\n", diff --git a/examples/GWMemory.ipynb b/examples/GWMemory.ipynb index 8b19b1d..7ec615c 100644 --- a/examples/GWMemory.ipynb +++ b/examples/GWMemory.ipynb @@ -54,45 +54,45 @@ "source": [ "fig = plt.figure(figsize=(12, 6))\n", "\n", - "qs = [1., 2.]\n", - "S1s = [[0., 0., 0.], [0.8, 0., 0.]]\n", - "S2s = [[0., 0., 0.], [0., 0.8, 0.]]\n", + "qs = [1.0, 2.0]\n", + "S1s = [[0.0, 0.0, 0.0], [0.8, 0.0, 0.0]]\n", + "S2s = [[0.0, 0.0, 0.0], [0.0, 0.8, 0.0]]\n", "\n", - "colours = ['r', 'b', 'g', 'k']\n", + "colours = [\"r\", \"b\", \"g\", \"k\"]\n", "\n", "times = np.linspace(-0.08, 0.02, 10001)\n", "\n", - "labels = ['Equal-mass, non-spinning']\n", + "labels = [\"Equal-mass, non-spinning\"]\n", "\n", "ax = [fig.add_subplot(2, 1, 1), fig.add_subplot(2, 1, 2)]\n", "\n", "parameters = dict(\n", - " total_mass=60, distance=400, model='NRSur7dq2', inc=np.pi/2, phase=0, times=times\n", + " total_mass=60, distance=400, model=\"NRSur7dq2\", inc=np.pi / 2, phase=0, times=times\n", ")\n", "\n", "for ii, q in enumerate(qs):\n", " parameters[\"q\"] = q\n", - " for jj, (S1, S2) in enumerate(zip(S1s, S2s)):\n", + " for jj, (S1, S2) in enumerate(zip(S1s, S2s)):\n", " colour = colours[ii + jj * 2]\n", - " \n", + "\n", " parameters[\"spin_1\"] = S1\n", " parameters[\"spin_2\"] = S2\n", - " \n", + "\n", " h_mem, times = gwmemory.time_domain_memory(**parameters)\n", "\n", - " ax[0].plot(times, h_mem['plus'] * 1e22, linestyle='-', color=colour)\n", - " ax[1].plot(times, h_mem['cross'] * 1e22, linestyle='-', color=colour)\n", + " ax[0].plot(times, h_mem[\"plus\"] * 1e22, linestyle=\"-\", color=colour)\n", + " ax[1].plot(times, h_mem[\"cross\"] * 1e22, linestyle=\"-\", color=colour)\n", "\n", " h_mem, times = gwmemory.time_domain_memory(\n", " **parameters, Lmax=2, modes=[(2, 2), (2, -2)]\n", " )\n", "\n", - " ax[0].plot(times, h_mem['plus'] * 1e22, linestyle=':', color=colour)\n", - " ax[1].plot(times, h_mem['cross'] * 1e22, linestyle=':', color=colour)\n", + " ax[0].plot(times, h_mem[\"plus\"] * 1e22, linestyle=\":\", color=colour)\n", + " ax[1].plot(times, h_mem[\"cross\"] * 1e22, linestyle=\":\", color=colour)\n", "\n", - "plt.xlabel('$t (s)$')\n", - "ax[0].set_ylabel('$\\delta h_{+} \\, \\left[10^{-22}\\\\right]$')\n", - "ax[1].set_ylabel('$\\delta h_{\\\\times} \\, \\left[10^{-22}\\\\right]$')\n", + "plt.xlabel(\"$t (s)$\")\n", + "ax[0].set_ylabel(\"$\\delta h_{+} \\, \\left[10^{-22}\\\\right]$\")\n", + "ax[1].set_ylabel(\"$\\delta h_{\\\\times} \\, \\left[10^{-22}\\\\right]$\")\n", "\n", "ax[0].set_xlim(-0.02, 0.02)\n", "ax[1].set_xlim(-0.02, 0.02)\n", @@ -117,24 +117,22 @@ "metadata": {}, "outputs": [], "source": [ - "qs = [1., 2.]\n", - "S1s = [[0., 0., 0.], [0.8, 0., 0.]]\n", - "S2s = [[0., 0., 0.], [0., 0.8, 0.]]\n", + "qs = [1.0, 2.0]\n", + "S1s = [[0.0, 0.0, 0.0], [0.8, 0.0, 0.0]]\n", + "S2s = [[0.0, 0.0, 0.0], [0.0, 0.8, 0.0]]\n", "\n", - "colours = ['r', 'b', 'g', 'k']\n", + "colours = [\"r\", \"b\", \"g\", \"k\"]\n", "\n", "mems_plus = {}\n", "mems_cross = {}\n", "\n", "times = np.linspace(-0.08, 0.02, 10001)\n", "\n", - "parameters = dict(\n", - " total_mass=60, distance=400, times=times\n", - ")\n", + "parameters = dict(total_mass=60, distance=400, times=times)\n", "\n", "for ii, q in enumerate(qs):\n", " parameters[\"q\"] = q\n", - " for jj, (S1, S2) in enumerate(zip(S1s, S2s)):\n", + " for jj, (S1, S2) in enumerate(zip(S1s, S2s)):\n", " parameters[\"spin_1\"] = S1\n", " parameters[\"spin_2\"] = S2\n", "\n", @@ -143,53 +141,63 @@ " hmem = {}\n", "\n", " modes = []\n", - " for l in range(2,5):\n", - " for m in np.flipud(range(0,l+1)):\n", - " modes += list(set([(l,m),(l,-m)]))\n", - " hmem[f'{l}{m}'], times = surr.time_domain_memory(inc=np.pi/2, phase=0, modes=modes)\n", + " for l in range(2, 5):\n", + " for m in np.flipud(range(0, l + 1)):\n", + " modes += list(set([(l, m), (l, -m)]))\n", + " hmem[f\"{l}{m}\"], times = surr.time_domain_memory(\n", + " inc=np.pi / 2, phase=0, modes=modes\n", + " )\n", "\n", " max_h_mem_plus = []\n", " max_h_mem_cross = []\n", " keys = []\n", - " for l in range(2,5):\n", - " for m in np.flipud(range(0,l+1)):\n", - " keys.append(f'{l}{m}')\n", - " max_h_mem_plus.append(hmem[f'{l}{m}']['plus'][-1])\n", - " max_h_mem_cross.append(hmem[f'{l}{m}']['cross'][-1])\n", + " for l in range(2, 5):\n", + " for m in np.flipud(range(0, l + 1)):\n", + " keys.append(f\"{l}{m}\")\n", + " max_h_mem_plus.append(hmem[f\"{l}{m}\"][\"plus\"][-1])\n", + " max_h_mem_cross.append(hmem[f\"{l}{m}\"][\"cross\"][-1])\n", "\n", - " mems_plus[f'{q}{S1}{S2}'] = np.array(max_h_mem_plus)\n", - " mems_cross[f'{q}{S1}{S2}'] = np.array(max_h_mem_cross)\n", + " mems_plus[f\"{q}{S1}{S2}\"] = np.array(max_h_mem_plus)\n", + " mems_cross[f\"{q}{S1}{S2}\"] = np.array(max_h_mem_cross)\n", "\n", - "colours = ['r', 'b', 'g', 'k']\n", + "colours = [\"r\", \"b\", \"g\", \"k\"]\n", "\n", - "spin_keys = [f'{S1}{S2}' for S1, S2 in zip(S1s, S2s)]\n", + "spin_keys = [f\"{S1}{S2}\" for S1, S2 in zip(S1s, S2s)]\n", "\n", "fig = plt.figure(figsize=(12, 4))\n", "\n", "for ii, q in enumerate(qs):\n", - " for jj, (S1, S2) in enumerate(zip(S1s, S2s)):\n", + " for jj, (S1, S2) in enumerate(zip(S1s, S2s)):\n", "\n", - " key = f'{q}{S1}{S2}'\n", + " key = f\"{q}{S1}{S2}\"\n", " max_h_mems_plus = mems_plus[key]\n", " max_h_mems_cross = mems_cross[key]\n", "\n", " plt.plot(\n", - " range(len(max_h_mems_plus)), max_h_mems_plus * 1e22, color=colours[ii + jj * 2], linestyle='-',\n", - " label=f'$q={q}$, $S_1={S1}$, $S_2={S2}$'\n", + " range(len(max_h_mems_plus)),\n", + " max_h_mems_plus * 1e22,\n", + " color=colours[ii + jj * 2],\n", + " linestyle=\"-\",\n", + " label=f\"$q={q}$, $S_1={S1}$, $S_2={S2}$\",\n", + " )\n", + " plt.plot(\n", + " range(len(max_h_mems_cross)),\n", + " max_h_mems_cross * 1e22,\n", + " color=colours[ii + jj * 2],\n", + " linestyle=\"--\",\n", " )\n", - " plt.plot(range(len(max_h_mems_cross)), max_h_mems_cross * 1e22, color=colours[ii + jj * 2], linestyle='--')\n", "\n", "plt.xticks(range(len(max_h_mems_plus)))\n", - "keys = [f'({ell},{delta_m})' for ell, delta_m in keys]\n", + "keys = [f\"({ell},{delta_m})\" for ell, delta_m in keys]\n", "ax = plt.gca()\n", "ax.set_xticklabels(keys)\n", - "plt.xlim(0, len(max_h_mems_plus)-1)\n", - "ax.grid(axis='both')\n", + "plt.xlim(0, len(max_h_mems_plus) - 1)\n", + "ax.grid(axis=\"both\")\n", "\n", - "plt.xlabel('Oscillatory $(\\ell, |m|)_{\\mathrm{last}}$')\n", - "plt.ylabel('$\\delta h(t=\\infty) \\, \\left[10^{-22}\\\\right]$')\n", + "plt.xlabel(\"Oscillatory $(\\ell, |m|)_{\\mathrm{last}}$\")\n", + "plt.ylabel(\"$\\delta h(t=\\infty) \\, \\left[10^{-22}\\\\right]$\")\n", "plt.legend(loc=\"lower right\")\n", - " \n", + "\n", "plt.tight_layout()\n", "plt.show()\n", "plt.close()" @@ -212,24 +220,22 @@ "source": [ "fig = plt.figure(figsize=(12, 4))\n", "\n", - "qs = [1., 2.]\n", - "S1s = [[0., 0., 0.], [0.8, 0., 0.]]\n", - "S2s = [[0., 0., 0.], [0., 0.8, 0.]]\n", + "qs = [1.0, 2.0]\n", + "S1s = [[0.0, 0.0, 0.0], [0.8, 0.0, 0.0]]\n", + "S2s = [[0.0, 0.0, 0.0], [0.0, 0.8, 0.0]]\n", "\n", - "colours = ['r', 'b', 'g', 'k']\n", + "colours = [\"r\", \"b\", \"g\", \"k\"]\n", "\n", "times = np.linspace(-0.08, 0.02, 10001)\n", "\n", - "parameters = dict(\n", - " total_mass=60, distance=400, times=times, model=\"NRSur7dq2\"\n", - ")\n", + "parameters = dict(total_mass=60, distance=400, times=times, model=\"NRSur7dq2\")\n", "\n", "for ii, q in enumerate(qs):\n", " parameters[\"q\"] = q\n", - " for jj, (S1, S2) in enumerate(zip(S1s, S2s)):\n", + " for jj, (S1, S2) in enumerate(zip(S1s, S2s)):\n", " parameters[\"spin_1\"] = S1\n", " parameters[\"spin_2\"] = S2\n", - " \n", + "\n", " h_mem, times = gwmemory.time_domain_memory(**parameters)\n", "\n", " colour = colours[ii + jj * 2]\n", @@ -238,18 +244,26 @@ " if ell <= 4:\n", " if ell == 2 and delta_m == 0:\n", " plt.scatter(\n", - " ell**2 + delta_m + ell, abs(h_mem[(ell, delta_m)][-1]), marker='x',\n", - " s=200, color=colour, label=f'$q={q}$, $S_1={S1}$, $S_2={S2}$'\n", + " ell**2 + delta_m + ell,\n", + " abs(h_mem[(ell, delta_m)][-1]),\n", + " marker=\"x\",\n", + " s=200,\n", + " color=colour,\n", + " label=f\"$q={q}$, $S_1={S1}$, $S_2={S2}$\",\n", " )\n", " else:\n", " plt.scatter(\n", - " ell**2 + delta_m + ell, abs(h_mem[(ell, delta_m)][-1]), marker='x',\n", - " s=200, color=colour\n", + " ell**2 + delta_m + ell,\n", + " abs(h_mem[(ell, delta_m)][-1]),\n", + " marker=\"x\",\n", + " s=200,\n", + " color=colour,\n", " )\n", " plt.plot(\n", " [ell**2 + delta_m + ell, ell**2 + delta_m + ell],\n", " [0, abs(h_mem[(ell, delta_m)][-1])],\n", - " color=colour, alpha=0.2\n", + " color=colour,\n", + " alpha=0.2,\n", " )\n", "\n", "plt.xlim(3, 5**2)\n", @@ -257,16 +271,16 @@ "keys = [f\"({ell, delta_m})\" for ell, delta_m in gwmemory.harmonics.lmax_modes(4)]\n", "ax = plt.gca()\n", "ax.set_xticklabels(keys)\n", - "plt.xlabel('Memory $(\\ell, m)$')\n", + "plt.xlabel(\"Memory $(\\ell, m)$\")\n", "\n", "plt.ylim(1e-27, 1e-21)\n", - "plt.yscale('log')\n", - "plt.ylabel('$|\\delta h_{\\ell m}(t=\\infty)|$')\n", + "plt.yscale(\"log\")\n", + "plt.ylabel(\"$|\\delta h_{\\ell m}(t=\\infty)|$\")\n", "\n", "ax = plt.gca()\n", "ax.set_yticks(np.logspace(-27, -21, 7))\n", - "ax.grid(axis='y')\n", - " \n", + "ax.grid(axis=\"y\")\n", + "\n", "plt.tight_layout()\n", "plt.show()\n", "plt.close()" @@ -298,52 +312,83 @@ "\n", "fig = plt.figure(figsize=(24, 10))\n", "\n", - "qs = [1.]\n", - "S1s = [[0., 0., 0.], [0.8, 0., 0.]]\n", - "S2s = [[0., 0., 0.], [0., 0.8, 0.]]\n", + "qs = [1.0]\n", + "S1s = [[0.0, 0.0, 0.0], [0.8, 0.0, 0.0]]\n", + "S2s = [[0.0, 0.0, 0.0], [0.0, 0.8, 0.0]]\n", "\n", - "labels = ['Non-spinning', 'Precessing']\n", + "labels = [\"Non-spinning\", \"Precessing\"]\n", "\n", - "parameters = dict(\n", - " total_mass=60, distance=400, times=times, model=\"NRSur7dq2\"\n", - ")\n", + "parameters = dict(total_mass=60, distance=400, times=times, model=\"NRSur7dq2\")\n", "\n", "for ii, q in enumerate(qs):\n", " parameters[\"q\"] = q\n", - " for jj, (S1, S2) in enumerate(zip(S1s, S2s)):\n", + " for jj, (S1, S2) in enumerate(zip(S1s, S2s)):\n", " parameters[\"spin_1\"] = S1\n", " parameters[\"spin_2\"] = S2\n", - " \n", + "\n", " h_mem_lm, times = gwmemory.time_domain_memory(**parameters)\n", "\n", - " inc_array = np.linspace(0, np.pi, 200) - np.pi/2\n", + " inc_array = np.linspace(0, np.pi, 200) - np.pi / 2\n", " pol_array = np.linspace(0, 2 * np.pi, 200) - np.pi\n", " pols, incs = np.meshgrid(pol_array, inc_array)\n", " pols_deg = pols * 180 / np.pi\n", " incs_deg = incs * 180 / np.pi\n", - " y_lm = {(ell, m): gwmemory.harmonics.sYlm(-2, ell, m, incs + np.pi/2, pols)\n", - " for ell, m in h_mem_lm.keys()}\n", + " y_lm = {\n", + " (ell, m): gwmemory.harmonics.sYlm(-2, ell, m, incs + np.pi / 2, pols)\n", + " for ell, m in h_mem_lm.keys()\n", + " }\n", + "\n", + " orientation_map = np.sum(\n", + " [y_lm[key] * h_mem_lm[key][-1] for key in y_lm], axis=0\n", + " )\n", "\n", - " orientation_map = np.sum([y_lm[key] * h_mem_lm[key][-1] for key in y_lm], axis=0)\n", - " \n", " ax = plt.subplot(2, 2, ii * 4 + jj * 2 + 1)\n", - " m = Basemap(projection='moll',lon_0=-180,resolution='c')\n", - " m.contourf(pols_deg, incs_deg, orientation_map.real, 100,\n", - " cmap=plt.cm.jet, latlon=True, levels=np.linspace(-2.5e-22, 2.5e-22, 201))\n", - "\n", - " if jj==0:\n", - " ax.annotate(text='$+$-polarized', xy=(0.08, 0.92), xycoords='axes fraction',\n", - " fontsize=30, horizontalalignment='center')\n", - " \n", + " m = Basemap(projection=\"moll\", lon_0=-180, resolution=\"c\")\n", + " m.contourf(\n", + " pols_deg,\n", + " incs_deg,\n", + " orientation_map.real,\n", + " 100,\n", + " cmap=plt.cm.jet,\n", + " latlon=True,\n", + " levels=np.linspace(-2.5e-22, 2.5e-22, 201),\n", + " )\n", + "\n", + " if jj == 0:\n", + " ax.annotate(\n", + " text=\"$+$-polarized\",\n", + " xy=(0.08, 0.92),\n", + " xycoords=\"axes fraction\",\n", + " fontsize=30,\n", + " horizontalalignment=\"center\",\n", + " )\n", + "\n", " ax = plt.subplot(2, 2, ii * 4 + jj * 2 + 2)\n", - " m = Basemap(projection='moll',lon_0=-180,resolution='c')\n", - " m.contourf(pols_deg, incs_deg, orientation_map.imag, 100,\n", - " cmap=plt.cm.jet, latlon=True, levels=np.linspace(-2.5e-22, 2.5e-22, 201))\n", - " if jj==0:\n", - " ax.annotate(text='$\\\\times$-polarized', xy=(0.08, 0.92), xycoords='axes fraction',\n", - " fontsize=30, horizontalalignment='center')\n", - " ax.annotate(text=labels[jj], xy=(0.9, 0.92), xycoords='axes fraction',\n", - " fontsize=30, horizontalalignment='center')\n", + " m = Basemap(projection=\"moll\", lon_0=-180, resolution=\"c\")\n", + " m.contourf(\n", + " pols_deg,\n", + " incs_deg,\n", + " orientation_map.imag,\n", + " 100,\n", + " cmap=plt.cm.jet,\n", + " latlon=True,\n", + " levels=np.linspace(-2.5e-22, 2.5e-22, 201),\n", + " )\n", + " if jj == 0:\n", + " ax.annotate(\n", + " text=\"$\\\\times$-polarized\",\n", + " xy=(0.08, 0.92),\n", + " xycoords=\"axes fraction\",\n", + " fontsize=30,\n", + " horizontalalignment=\"center\",\n", + " )\n", + " ax.annotate(\n", + " text=labels[jj],\n", + " xy=(0.9, 0.92),\n", + " xycoords=\"axes fraction\",\n", + " fontsize=30,\n", + " horizontalalignment=\"center\",\n", + " )\n", "\n", "plt.tight_layout()\n", "plt.show()\n", @@ -367,57 +412,55 @@ "source": [ "fig = plt.figure(figsize=(12, 4))\n", "\n", - "qs = [1., 2.]\n", - "S1s = [[0., 0., 0.], [0.8, 0., 0.]]\n", - "S2s = [[0., 0., 0.], [0., 0.8, 0.]]\n", + "qs = [1.0, 2.0]\n", + "S1s = [[0.0, 0.0, 0.0], [0.8, 0.0, 0.0]]\n", + "S2s = [[0.0, 0.0, 0.0], [0.0, 0.8, 0.0]]\n", "\n", - "colours = ['r', 'b', 'g', 'k']\n", + "colours = [\"r\", \"b\", \"g\", \"k\"]\n", "\n", "times = np.linspace(-0.08, 0.02, 10001)\n", "\n", - "parameters = dict(\n", - " total_mass=60, distance=400, times=times\n", - ")\n", + "parameters = dict(total_mass=60, distance=400, times=times)\n", "\n", "for ii, q in enumerate(qs):\n", " parameters[\"q\"] = q\n", - " for jj, (S1, S2) in enumerate(zip(S1s, S2s)):\n", + " for jj, (S1, S2) in enumerate(zip(S1s, S2s)):\n", " parameters[\"spin_1\"] = S1\n", " parameters[\"spin_2\"] = S2\n", - " \n", + "\n", " surr = gwmemory.waveforms.Surrogate(**parameters)\n", " osc = gwmemory.utils.combine_modes(surr.h_lm, np.pi / 2, 0)\n", " modes = gwmemory.harmonics.lmax_modes(4)\n", " old_h_mem = np.array(0 * (1 + 1j))\n", - " delta_h = [max(abs(osc['plus'] - 1j * osc['cross']))]\n", + " delta_h = [max(abs(osc[\"plus\"] - 1j * osc[\"cross\"]))]\n", "\n", " for kk in range(1, 5):\n", " h_mem, times = surr.time_domain_memory(np.pi / 2, 0)\n", - " delta_h.append(h_mem['plus'][-1] - 1j * h_mem['cross'][-1] - old_h_mem)\n", - " old_h_mem = h_mem['plus'][-1] - 1j * h_mem['cross'][-1]\n", + " delta_h.append(h_mem[\"plus\"][-1] - 1j * h_mem[\"cross\"][-1] - old_h_mem)\n", + " old_h_mem = h_mem[\"plus\"][-1] - 1j * h_mem[\"cross\"][-1]\n", " h_mem_lm, times = surr.time_domain_memory()\n", " surr = gwmemory.waveforms.Surrogate(**parameters)\n", " for key in h_mem_lm:\n", " if key[0] <= 4:\n", " surr.h_lm[key] += h_mem_lm[key]\n", - " \n", + "\n", " colour = colours[ii + jj * 2]\n", "\n", " plt.semilogy(\n", " range(5),\n", " np.abs(delta_h),\n", - " marker='x',\n", - " label=f'$q={q}$, $S_1={S1}$, $S_2={S2}$',\n", + " marker=\"x\",\n", + " label=f\"$q={q}$, $S_1={S1}$, $S_2={S2}$\",\n", " color=colour,\n", " )\n", - " plt.xlabel('Memory order, $i$')\n", - " plt.ylabel('$\\max(|h^i|)$')\n", - " \n", + " plt.xlabel(\"Memory order, $i$\")\n", + " plt.ylabel(\"$\\max(|h^i|)$\")\n", + "\n", "plt.xticks(range(5))\n", "plt.legend(loc=\"lower left\")\n", "\n", "ax = plt.gca()\n", - "ax.grid(True, axis='y')\n", + "ax.grid(True, axis=\"y\")\n", "\n", "plt.tight_layout()\n", "plt.show()\n", @@ -456,26 +499,33 @@ "S1 = [0.0, 0.0, 0.0]\n", "S2 = [0.0, 0.0, 0.0]\n", "\n", - "colours = ['r', 'b', 'g', 'k']\n", + "colours = [\"r\", \"b\", \"g\", \"k\"]\n", "\n", "parameters = dict(\n", - " q=1, spin_1=S1, spin_2=S2,\n", - " total_mass=60, distance=400,\n", - " inc=np.pi / 2, phase=0.0\n", + " q=1, spin_1=S1, spin_2=S2, total_mass=60, distance=400, inc=np.pi / 2, phase=0.0\n", ")\n", "\n", "\n", - "for ii, model in enumerate(['NRSur7dq2', 'IMRPhenomD', 'SEOBNRv4', 'MWM']):\n", + "for ii, model in enumerate([\"NRSur7dq2\", \"IMRPhenomD\", \"SEOBNRv4\", \"MWM\"]):\n", " h_mem, times = gwmemory.time_domain_memory(**parameters, model=model)\n", "\n", - " plt.plot(times, h_mem['plus'] - h_mem['plus'][np.argmin(abs(times+0.25))],\n", - " linestyle='-', color=colours[ii], label=model)\n", - " plt.plot(times, h_mem['cross'] - h_mem['cross'][np.argmin(abs(times+0.25))],\n", - " linestyle='--', color=colours[ii])\n", - " \n", - "plt.xlabel('$t (s)$')\n", - "plt.ylabel('$\\delta h$')\n", - "plt.legend(loc='upper left', fontsize=20)\n", + " plt.plot(\n", + " times,\n", + " h_mem[\"plus\"] - h_mem[\"plus\"][np.argmin(abs(times + 0.25))],\n", + " linestyle=\"-\",\n", + " color=colours[ii],\n", + " label=model,\n", + " )\n", + " plt.plot(\n", + " times,\n", + " h_mem[\"cross\"] - h_mem[\"cross\"][np.argmin(abs(times + 0.25))],\n", + " linestyle=\"--\",\n", + " color=colours[ii],\n", + " )\n", + "\n", + "plt.xlabel(\"$t (s)$\")\n", + "plt.ylabel(\"$\\delta h$\")\n", + "plt.legend(loc=\"upper left\", fontsize=20)\n", "\n", "plt.xlim(-0.25, 0.02)\n", "\n", @@ -499,9 +549,9 @@ "metadata": {}, "outputs": [], "source": [ - "q = 1.\n", - "S1 = [0., 0., 0.]\n", - "S2 = [0., 0., 0.]\n", + "q = 1.0\n", + "S1 = [0.0, 0.0, 0.0]\n", + "S2 = [0.0, 0.0, 0.0]\n", "\n", "times = np.linspace(-0.08, 0.02, 10001)\n", "surr = gwmemory.waveforms.Surrogate(\n", @@ -516,27 +566,27 @@ "\n", "fig = plt.figure(figsize=(12, 6))\n", "fig.add_subplot(2, 1, 1)\n", - "plt.plot(times, oscillatory['plus'], linestyle='--', color='b', alpha=0.5)\n", - "plt.plot(times, oscillatory['cross'], linestyle='--', color='r', alpha=0.5)\n", - "plt.plot(times, memory['plus'], linestyle='-.', color='b', alpha=0.5)\n", - "plt.plot(times, memory['cross'], linestyle='-.', color='r', alpha=0.5)\n", - "plt.plot(times, oscillatory['plus'] + memory['plus'], color='b')\n", - "plt.plot(times, oscillatory['cross'] + memory['cross'], color='r')\n", + "plt.plot(times, oscillatory[\"plus\"], linestyle=\"--\", color=\"b\", alpha=0.5)\n", + "plt.plot(times, oscillatory[\"cross\"], linestyle=\"--\", color=\"r\", alpha=0.5)\n", + "plt.plot(times, memory[\"plus\"], linestyle=\"-.\", color=\"b\", alpha=0.5)\n", + "plt.plot(times, memory[\"cross\"], linestyle=\"-.\", color=\"r\", alpha=0.5)\n", + "plt.plot(times, oscillatory[\"plus\"] + memory[\"plus\"], color=\"b\")\n", + "plt.plot(times, oscillatory[\"cross\"] + memory[\"cross\"], color=\"r\")\n", "plt.xlabel(\"$t [s]$\")\n", "plt.ylabel(\"Strain\")\n", - "plt.axhline(0, linestyle=':', color='k')\n", + "plt.axhline(0, linestyle=\":\", color=\"k\")\n", "plt.xlim(-0.08, 0.0)\n", "\n", "fig.add_subplot(2, 1, 2)\n", - "plt.plot(times, oscillatory['plus'], linestyle='--', color='b', alpha=0.5)\n", - "plt.plot(times, oscillatory['cross'], linestyle='--', color='r', alpha=0.5)\n", - "plt.plot(times, memory['plus'], linestyle='-.', color='b', alpha=0.5)\n", - "plt.plot(times, memory['cross'], linestyle='-.', color='r', alpha=0.5)\n", - "plt.plot(times, oscillatory['plus'] + memory['plus'], color='b')\n", - "plt.plot(times, oscillatory['cross'] + memory['cross'], color='r')\n", + "plt.plot(times, oscillatory[\"plus\"], linestyle=\"--\", color=\"b\", alpha=0.5)\n", + "plt.plot(times, oscillatory[\"cross\"], linestyle=\"--\", color=\"r\", alpha=0.5)\n", + "plt.plot(times, memory[\"plus\"], linestyle=\"-.\", color=\"b\", alpha=0.5)\n", + "plt.plot(times, memory[\"cross\"], linestyle=\"-.\", color=\"r\", alpha=0.5)\n", + "plt.plot(times, oscillatory[\"plus\"] + memory[\"plus\"], color=\"b\")\n", + "plt.plot(times, oscillatory[\"cross\"] + memory[\"cross\"], color=\"r\")\n", "plt.xlabel(\"$t [s]$\")\n", "plt.ylabel(\"Strain\")\n", - "plt.axhline(0, linestyle=':', color='k')\n", + "plt.axhline(0, linestyle=\":\", color=\"k\")\n", "plt.xlim(-0.0, 0.02)\n", "\n", "plt.tight_layout()\n", @@ -562,30 +612,37 @@ "source": [ "fig = plt.figure(figsize=(12, 4))\n", "\n", - "q = 1.\n", + "q = 1.0\n", "S1 = [0, 0.8, 0]\n", "S2 = [0.8, 0, 0]\n", "\n", "times = np.linspace(-0.98, 0.02, 10000)\n", "\n", - "colours = ['r', 'b', 'g', 'k']\n", + "colours = [\"r\", \"b\", \"g\", \"k\"]\n", "\n", "h_mem, frequencies = gwmemory.frequency_domain_memory(\n", - " q=q, spin_1=S1, spin_2=S2, total_mass=60., distance=400.,\n", - " model='NRSur7dq2', inc=np.pi/2, phase=0., times=times\n", + " q=q,\n", + " spin_1=S1,\n", + " spin_2=S2,\n", + " total_mass=60.0,\n", + " distance=400.0,\n", + " model=\"NRSur7dq2\",\n", + " inc=np.pi / 2,\n", + " phase=0.0,\n", + " times=times,\n", ")\n", "\n", - "plt.loglog(frequencies, abs(h_mem['plus']), linestyle='-', color='r', label=model)\n", - "plt.loglog(frequencies, abs(h_mem['cross']), linestyle='--', color='r')\n", - " \n", - "plt.xlabel('$f$ (Hz)')\n", - "plt.ylabel('Strain [Hz$^{-1/2}$]')\n", + "plt.loglog(frequencies, abs(h_mem[\"plus\"]), linestyle=\"-\", color=\"r\", label=model)\n", + "plt.loglog(frequencies, abs(h_mem[\"cross\"]), linestyle=\"--\", color=\"r\")\n", + "\n", + "plt.xlabel(\"$f$ (Hz)\")\n", + "plt.ylabel(\"Strain [Hz$^{-1/2}$]\")\n", "\n", "plt.xlim(10, 2048)\n", "plt.ylim(1e-27, 1e-23)\n", "\n", "ax = plt.gca()\n", - "ax.grid(True, axis='y')\n", + "ax.grid(True, axis=\"y\")\n", "\n", "plt.tight_layout()\n", "plt.show()\n", diff --git a/gwmemory/__init__.py b/gwmemory/__init__.py index 658cb24..0b0349a 100644 --- a/gwmemory/__init__.py +++ b/gwmemory/__init__.py @@ -1,5 +1,5 @@ from . import angles, gwmemory, harmonics, qnms, utils, waveforms -from .gwmemory import time_domain_memory, frequency_domain_memory +from .gwmemory import frequency_domain_memory, time_domain_memory try: from ._version import version as __version__ diff --git a/gwmemory/angles.py b/gwmemory/angles.py index 9fcd314..89fca44 100644 --- a/gwmemory/angles.py +++ b/gwmemory/angles.py @@ -1,5 +1,6 @@ #!/usr/bin/python3 from functools import lru_cache +from typing import Tuple import numpy as np from sympy.physics.wigner import wigner_3j @@ -8,7 +9,7 @@ @lru_cache -def memory_correction(ell, ss=0): +def memory_correction(ell: int, ss: int = 0) -> float: """ Correction to the Gamma function for the operator in Eq. (12) of arXiv:2011.01309 @@ -28,13 +29,14 @@ def memory_correction(ell, ss=0): if ell < 2: return 0 return ( - ((ell - (ss - 1)) * (ell + ss) * (ell - (ss - 2)) * (ell + (ss - 1)))**0.5 - * 4 / ((ell + 2) * (ell + 1) * ell * (ell - 1)) + ((ell - (ss - 1)) * (ell + ss) * (ell - (ss - 2)) * (ell + (ss - 1))) ** 0.5 + * 4 + / ((ell + 2) * (ell + 1) * ell * (ell - 1)) ) @lru_cache() -def analytic_gamma(lm1, lm2, ell): +def analytic_gamma(lm1: Tuple[int, int], lm2: Tuple[int, int], ell: int) -> float: """ Analytic function to compute gamma_lmlm_l Eq. (8) of arXiv:1807.0090 @@ -61,14 +63,27 @@ def analytic_gamma(lm1, lm2, ell): m3 = m1 + m2 return ( (-1.0) ** (ell1 + ell2 + ell + m3 + m2) - * (2 * ell1 + 1)**0.5 * (2 * ell2 + 1)**0.5 * (2 * ell + 1)**0.5 - * float(wigner_3j(ell1, ell2, ell, s1, s2, -s3) * wigner_3j(ell1, ell2, ell, m1, m2, -m3)) - * np.pi**0.5 / 2 + * (2 * ell1 + 1) ** 0.5 + * (2 * ell2 + 1) ** 0.5 + * (2 * ell + 1) ** 0.5 + * float( + wigner_3j(ell1, ell2, ell, s1, s2, -s3) + * wigner_3j(ell1, ell2, ell, m1, m2, -m3) + ) + * np.pi**0.5 + / 2 * memory_correction(ell) ) -def gamma(lm1, lm2, incs=None, theta=None, phi=None, y_lmlm_factor=None): +def gamma( + lm1: str, + lm2: str, + incs: np.ndarray = None, + theta: np.ndarray = None, + phi: np.ndarray = None, + y_lmlm_factor: np.ndarray = None, +) -> list: """ Coefficients mapping the spherical harmonic components of the oscillatory strain to the memory. @@ -131,21 +146,23 @@ def gamma(lm1, lm2, incs=None, theta=None, phi=None, y_lmlm_factor=None): if ell < abs(delta_m): gammas.append(0) else: - gammas.append(np.real( - 2 - * np.pi - * np.trapz( - lambda_lm1_lm2 - * np.conjugate(harm[f"{ell}{-delta_m}"]) - * sin_inc, - incs, + gammas.append( + np.real( + 2 + * np.pi + * np.trapz( + lambda_lm1_lm2 + * np.conjugate(harm[f"{ell}{-delta_m}"]) + * sin_inc, + incs, + ) ) - )) + ) return gammas -def ylmlm_factor(theta, phi, lm1, lm2): +def ylmlm_factor(theta: np.ndarray, phi: np.ndarray, lm1: str, lm2: str) -> np.ndarray: if theta is None: theta = np.linspace(0, np.pi, 250) if phi is None: @@ -166,7 +183,15 @@ def ylmlm_factor(theta, phi, lm1, lm2): return y_lmlm_factor, theta, phi -def lambda_matrix(inc, phase, lm1, lm2, theta=None, phi=None, y_lmlm_factor=None): +def lambda_matrix( + inc: float, + phase: float, + lm1: str, + lm2: str, + theta: np.ndarray = None, + phi: np.ndarray = None, + y_lmlm_factor: np.ndarray = None, +) -> np.ndarray: r""" Angular integral for a specific ll'mm' as given by equation 7 of Talbot et al. (2018), arXiv:1807.00990. @@ -205,17 +230,21 @@ def lambda_matrix(inc, phase, lm1, lm2, theta=None, phi=None, y_lmlm_factor=None if y_lmlm_factor is None: y_lmlm_factor, theta, phi = ylmlm_factor(theta=theta, phi=phi, lm1=lm1, lm2=lm2) - n = np.array([ - np.outer(np.cos(phi), np.sin(theta)), - np.outer(np.sin(phi), np.sin(theta)), - np.outer(np.ones_like(phi), np.cos(theta)), - ]) - line_of_sight = np.array([np.sin(inc) * np.cos(phase), np.sin(inc) * np.sin(phase), np.cos(inc)]) + n = np.array( + [ + np.outer(np.cos(phi), np.sin(theta)), + np.outer(np.sin(phi), np.sin(theta)), + np.outer(np.ones_like(phi), np.cos(theta)), + ] + ) + line_of_sight = np.array( + [np.sin(inc) * np.cos(phase), np.sin(inc) * np.sin(phase), np.cos(inc)] + ) n_dot_line_of_sight = sum(n_i * N_i for n_i, N_i in zip(n, line_of_sight)) n_dot_line_of_sight[n_dot_line_of_sight == 1] = 0 denominator = 1 / (1 - n_dot_line_of_sight) - sin_array = np.outer(phi ** 0, np.sin(theta)) + sin_array = np.outer(phi**0, np.sin(theta)) angle_integrals_r = np.zeros((3, 3)) angle_integrals_i = np.zeros((3, 3)) @@ -228,8 +257,9 @@ def lambda_matrix(inc, phase, lm1, lm2, theta=None, phi=None, y_lmlm_factor=None * y_lmlm_factor * ( n[j] * n[k] - - (n[j] * line_of_sight[k] + n[k] * line_of_sight[j]) * n_dot_line_of_sight - + line_of_sight[j] * line_of_sight[k] * n_dot_line_of_sight ** 2 + - (n[j] * line_of_sight[k] + n[k] * line_of_sight[j]) + * n_dot_line_of_sight + + line_of_sight[j] * line_of_sight[k] * n_dot_line_of_sight**2 ) ) angle_integrals_r[j, k] = np.trapz(np.trapz(np.real(integrand), theta), phi) @@ -244,7 +274,15 @@ def lambda_matrix(inc, phase, lm1, lm2, theta=None, phi=None, y_lmlm_factor=None return lambda_mat -def lambda_lmlm(inc, phase, lm1, lm2, theta=None, phi=None, y_lmlm_factor=None): +def lambda_lmlm( + inc: float, + phase: float, + lm1: str, + lm2: str, + theta: np.ndarray = None, + phi: np.ndarray = None, + y_lmlm_factor: np.ndarray = None, +) -> complex: r""" Angular integral for a specific ll'mm' as given by equation 7 of Talbot et al. (2018), arXiv:1807.00990. @@ -290,7 +328,9 @@ def lambda_lmlm(inc, phase, lm1, lm2, theta=None, phi=None, y_lmlm_factor=None): return lambda_lmlm -def omega_ij_to_omega_pol(omega_ij, inc, phase): +def omega_ij_to_omega_pol( + omega_ij: np.ndarray, inc: float, phase: float +) -> Tuple[np.ndarray, np.ndarray]: """ Map from strain tensor to plus and cross modes. @@ -322,7 +362,7 @@ def omega_ij_to_omega_pol(omega_ij, inc, phase): return omega_plus, omega_cross -def plus_tensor(wx, wy): +def plus_tensor(wx: np.ndarray, wy: np.ndarray) -> np.ndarray: """ Calculate the plus polarization tensor for some basis. c.f., eq. 2 of https://arxiv.org/pdf/1710.03794.pdf @@ -331,7 +371,7 @@ def plus_tensor(wx, wy): return e_plus -def cross_tensor(wx, wy): +def cross_tensor(wx: np.ndarray, wy: np.ndarray) -> np.ndarray: """ Calculate the cross polarization tensor for some basis. c.f., eq. 2 of https://arxiv.org/pdf/1710.03794.pdf @@ -340,7 +380,9 @@ def cross_tensor(wx, wy): return e_cross -def wave_frame(theta, phi, psi=0): +def wave_frame( + theta: float, phi: float, psi: float = 0 +) -> Tuple[np.ndarray, np.ndarray]: """ Generate wave-frame basis from three angles, see Nishizawa et al. (2009) """ diff --git a/gwmemory/gwmemory.py b/gwmemory/gwmemory.py index a5c2545..919f067 100644 --- a/gwmemory/gwmemory.py +++ b/gwmemory/gwmemory.py @@ -1,21 +1,24 @@ -from . import waveforms -from . import utils import inspect +from typing import Tuple + +import numpy as np + +from . import utils, waveforms def time_domain_memory( - model=None, - h_lm=None, - times=None, - q=None, - total_mass=None, - spin_1=None, - spin_2=None, - distance=None, - inc=None, - phase=None, + model: str = None, + h_lm: dict = None, + times: np.ndarray = None, + q: float = None, + total_mass: float = None, + spin_1: Tuple[float, float, float] = None, + spin_2: Tuple[float, float, float] = None, + distance: float = None, + inc: float = None, + phase: float = None, **kwargs, -): +) -> Tuple[dict, np.ndarray]: """ Calculate the time domain memory waveform according to __reference__. @@ -137,16 +140,16 @@ def time_domain_memory( def frequency_domain_memory( - model=None, - q=None, - total_mass=None, - spin_1=None, - spin_2=None, - distance=None, - inc=None, - phase=None, + model: str = None, + q: float = None, + total_mass: float = None, + spin_1: Tuple[float, float, float] = None, + spin_2: Tuple[float, float, float] = None, + distance: float = None, + inc: float = None, + phase: float = None, **kwargs, -): +) -> Tuple[dict, np.ndarray]: """ Calculate the frequency domain memory waveform according to __reference__. diff --git a/gwmemory/harmonics.py b/gwmemory/harmonics.py index 2ce7fd2..8277f79 100644 --- a/gwmemory/harmonics.py +++ b/gwmemory/harmonics.py @@ -9,24 +9,29 @@ # ---------------------------------------------------------- from functools import lru_cache from math import factorial +from typing import Union import numpy as np # coefficient function @lru_cache -def Cslm(s, l, m): +def Cslm(s: int, l: int, m: int) -> float: return np.sqrt(l * l * (4.0 * l * l - 1.0) / ((l * l - m * m) * (l * l - s * s))) # recursion function -def s_lambda_lm(s, l, m, x): +def s_lambda_lm( + s: int, l: int, m: int, x: Union[float, np.ndarray] +) -> Union[float, np.ndarray]: Pm = pow(-0.5, m) if m != s: Pm = Pm * pow(1.0 + x, (m - s) * 1.0 / 2) if m != -s: Pm = Pm * pow(1.0 - x, (m + s) * 1.0 / 2) - Pm = Pm * np.sqrt(factorial(2 * m + 1) * 1.0 / (4.0 * np.pi * factorial(m + s) * factorial(m - s))) + Pm = Pm * np.sqrt( + factorial(2 * m + 1) * 1.0 / (4.0 * np.pi * factorial(m + s) * factorial(m - s)) + ) if l == m: return Pm Pm1 = (x + s * 1.0 / (m + 1)) * Cslm(s, m + 1, m) * Pm @@ -42,7 +47,13 @@ def s_lambda_lm(s, l, m, x): return Pn -def sYlm(ss, ll, mm, theta, phi): +def sYlm( + ss: int, + ll: int, + mm: int, + theta: Union[float, np.ndarray], + phi: Union[float, np.ndarray], +) -> Union[float, np.ndarray]: """Calculate spin-weighted harmonic""" Pm = 1.0 l = ll @@ -67,6 +78,6 @@ def sYlm(ss, ll, mm, theta, phi): @lru_cache -def lmax_modes(lmax): +def lmax_modes(lmax: int) -> list: """Compute all (l, m) pairs with 2<=l<=lmax""" return [(l, m) for l in range(2, lmax + 1) for m in range(-l, l + 1)] diff --git a/gwmemory/qnms.py b/gwmemory/qnms.py index d0e033a..a309d86 100644 --- a/gwmemory/qnms.py +++ b/gwmemory/qnms.py @@ -2,11 +2,12 @@ # Modified Colm Talbot # This code package includes code for calculating the properties of quasinormal # black hole modes +from typing import Tuple import numpy as np -def freq_damping(mass, spin, nn=0): +def freq_damping(mass: float, spin: float, nn: int = 0) -> Tuple[float, float]: """ Calculate the quasinormal mode freq and damping time for a black hole. This version uses OBSERVER'S UNITS. @@ -20,10 +21,6 @@ def freq_damping(mass, spin, nn=0): BH mass in solar masses spin: float BH dimensionless spin - ell: int, optional - Spherical harmonic l, default=2 - mm: int, optional - Spherical harmonic m, default=2 nn: int, optional QNM harmonic, default=0 is the fundamental @@ -58,7 +55,7 @@ def freq_damping(mass, spin, nn=0): return omega_lmn, tau_lmn -def final_mass_spin(mass_1, mass_2): +def final_mass_spin(mass_1: float, mass_2: float) -> Tuple[float, float]: """ Given initial total mass and mass ratio, calculate the final mass, M_f and dimensionless spin parameter, jj, @@ -88,14 +85,14 @@ def final_mass_spin(mass_1, mass_2): # expansion coefficients for the mass m_f1 = 1.0 m_f2 = (np.sqrt(8.0 / 9.0) - 1.0) * eta - m_f3 = -0.498 * eta ** 2 + m_f3 = -0.498 * eta**2 # final mass final_mass = total_mass * (m_f1 + m_f2 + m_f3) # expansion coefficients for spin parameter a_f1 = np.sqrt(12.0) * eta - a_f2 = -2.9 * eta ** 2 + a_f2 = -2.9 * eta**2 # final spin parameter -- dimensions of mass a_f = final_mass * (a_f1 + a_f2) diff --git a/gwmemory/utils.py b/gwmemory/utils.py index 7981a8a..9ba3d57 100644 --- a/gwmemory/utils.py +++ b/gwmemory/utils.py @@ -1,21 +1,22 @@ -import numpy as np +from typing import Tuple -from .harmonics import sYlm, lmax_modes +import numpy as np +from .harmonics import lmax_modes, sYlm # Constants for conversions # taken from astropy==5.0.1 CC = 299792458.0 GG = 6.6743e-11 -SOLAR_MASS = 1.988409870698051e+30 +SOLAR_MASS = 1.988409870698051e30 KG = 1 / SOLAR_MASS -METRE = CC ** 2 / (GG * SOLAR_MASS) +METRE = CC**2 / (GG * SOLAR_MASS) SECOND = CC * METRE -MPC = 3.085677581491367e+22 +MPC = 3.085677581491367e22 -def nfft(ht, sampling_frequency): +def nfft(ht: np.ndarray, sampling_frequency: float) -> Tuple[np.ndarray, np.ndarray]: """ performs an FFT while keeping track of the frequency bins assumes input time series is real (positive frequencies only) @@ -52,7 +53,9 @@ def nfft(ht, sampling_frequency): return hf, ff -def load_sxs_waveform(file_name, modes=None, extraction="OutermostExtraction.dir"): +def load_sxs_waveform( + file_name: str, modes: list = None, extraction: str = "OutermostExtraction.dir" +) -> Tuple[dict, np.ndarray]: """ Load the spherical harmonic modes of an SXS numerical relativity waveform. @@ -86,7 +89,7 @@ def load_sxs_waveform(file_name, modes=None, extraction="OutermostExtraction.dir return output, times -def combine_modes(h_lm, inc, phase): +def combine_modes(h_lm: dict, inc: float, phase: float) -> np.ndarray: """ Calculate the plus and cross polarisations of the waveform from the spherical harmonic decomposition. diff --git a/gwmemory/waveforms/__init__.py b/gwmemory/waveforms/__init__.py index a2b9035..ce4a8cb 100644 --- a/gwmemory/waveforms/__init__.py +++ b/gwmemory/waveforms/__init__.py @@ -1,4 +1,4 @@ -from .base import MemoryGenerator +from .base import MemoryGenerator # isort: skip from . import approximant, mwm, nr, surrogate from .approximant import Approximant from .mwm import MWM diff --git a/gwmemory/waveforms/approximant.py b/gwmemory/waveforms/approximant.py index e75d9bc..e61958c 100644 --- a/gwmemory/waveforms/approximant.py +++ b/gwmemory/waveforms/approximant.py @@ -1,7 +1,9 @@ +from typing import Tuple + import numpy as np from ..harmonics import sYlm -from ..utils import combine_modes, MPC, SOLAR_MASS +from ..utils import MPC, SOLAR_MASS, combine_modes from . import MemoryGenerator @@ -13,18 +15,18 @@ class Approximant(MemoryGenerator): def __init__( self, - name, - q, - total_mass=60, - spin_1=None, - spin_2=None, - distance=400, - times=None, - minimum_frequency=20.0, - reference_frequency=20.0, - duration=4.0, - sampling_frequency=2048.0, - l_max=4, + name: str, + q: float, + total_mass: float = 60, + spin_1: Tuple[float, float, float] = None, + spin_2: Tuple[float, float, float] = None, + distance: float = 400, + times: np.ndarray = None, + minimum_frequency: float = 20.0, + reference_frequency: float = 20.0, + duration: float = 4.0, + sampling_frequency: float = 2048.0, + l_max: int = 4, ): """ Initialise Surrogate MemoryGenerator @@ -109,7 +111,7 @@ def name(self): return self._name @name.setter - def name(self, name): + def name(self, name: str): if name in self.no_hm: self._kind = "no_hm" elif name in self.td: @@ -123,7 +125,13 @@ def name(self, name): ) self._name = name - def time_domain_oscillatory(self, delta_t=None, modes=None, inc=None, phase=None): + def time_domain_oscillatory( + self, + delta_t: float = None, + modes: list = None, + inc: float = None, + phase: float = None, + ) -> Tuple[dict, np.ndarray]: """ Get the mode decomposition of the waveform approximant. @@ -152,16 +160,8 @@ def time_domain_oscillatory(self, delta_t=None, modes=None, inc=None, phase=None times: np.array Times on which waveform is evaluated. """ + import lalsimulation as lalsim from lal import CreateDict - from lalsimulation import ( - GetApproximantFromString, - SimInspiralTD, - SimInspiralChooseTDModes, - SimInspiralChooseFDModes, - SimInspiralCreateModeArray, - SimInspiralModeArrayActivateMode, - SimInspiralWaveformParamsInsertModeArray, - ) if self.h_lm is None: @@ -169,7 +169,7 @@ def time_domain_oscillatory(self, delta_t=None, modes=None, inc=None, phase=None f_min=self.minimum_frequency, f_ref=self.reference_frequency, phiRef=0.0, - approximant=GetApproximantFromString(self.name), + approximant=lalsim.GetApproximantFromString(self.name), LALpars=None, m1=self.m1_SI, m2=self.m2_SI, @@ -191,10 +191,10 @@ def time_domain_oscillatory(self, delta_t=None, modes=None, inc=None, phase=None if self._kind in ["td", "fd"]: if modes is not None: WFdict = CreateDict() - mode_array = SimInspiralCreateModeArray() + mode_array = lalsim.SimInspiralCreateModeArray() for mode in modes: - SimInspiralModeArrayActivateMode(mode_array, *mode) - SimInspiralWaveformParamsInsertModeArray(WFdict, modes) + lalsim.SimInspiralModeArrayActivateMode(mode_array, *mode) + lalsim.SimInspiralWaveformParamsInsertModeArray(WFdict, modes) else: WFdict = None params["LALpars"] = WFdict @@ -206,16 +206,20 @@ def time_domain_oscillatory(self, delta_t=None, modes=None, inc=None, phase=None duration = self.duration params["deltaF"] = 1 / duration params["f_max"] = self.sampling_frequency / 2 - waveform_modes = SimInspiralChooseFDModes(**params) + waveform_modes = lalsim.SimInspiralChooseFDModes(**params) times = np.arange(0, duration, 1 / self.sampling_frequency) h_lm = dict() while waveform_modes is not None: mode = (waveform_modes.l, waveform_modes.m) data = waveform_modes.mode.data.data[:-1] - h_lm[mode] = np.roll( - np.fft.ifft(np.roll(data, int(duration * params["f_max"]))), - int((duration - 1) * self.sampling_frequency)) * self.sampling_frequency + h_lm[mode] = ( + np.roll( + np.fft.ifft(np.roll(data, int(duration * params["f_max"]))), + int((duration - 1) * self.sampling_frequency), + ) + * self.sampling_frequency + ) if mode == (2, 2): times -= times[np.argmax(abs(h_lm[mode]))] waveform_modes = waveform_modes.next @@ -223,7 +227,7 @@ def time_domain_oscillatory(self, delta_t=None, modes=None, inc=None, phase=None del params["inclination"] params["r"] = params.pop("distance") params["lmax"] = self.l_max - waveform_modes = SimInspiralChooseTDModes(**params) + waveform_modes = lalsim.SimInspiralChooseTDModes(**params) times = np.arange(len(waveform_modes.mode.data.data)) * params["deltaT"] h_lm = dict() @@ -238,7 +242,7 @@ def time_domain_oscillatory(self, delta_t=None, modes=None, inc=None, phase=None params["eccentricity"] = 0.0 params["meanPerAno"] = 0.0 params["LALparams"] = params.pop("LALpars") - hplus, hcross = SimInspiralTD(**params) + hplus, hcross = lalsim.SimInspiralTD(**params) h = hplus.data.data - 1j * hcross.data.data h_22 = h / sYlm(-2, 2, 2, 0, 0) diff --git a/gwmemory/waveforms/base.py b/gwmemory/waveforms/base.py index fbf5e06..07416fc 100644 --- a/gwmemory/waveforms/base.py +++ b/gwmemory/waveforms/base.py @@ -1,8 +1,9 @@ import warnings +from typing import Tuple import numpy as np -from scipy.interpolate import interp1d from scipy.integrate import cumulative_trapezoid +from scipy.interpolate import interp1d from ..angles import analytic_gamma from ..harmonics import lmax_modes @@ -10,7 +11,7 @@ class MemoryGenerator(object): - def __init__(self, name, h_lm, times, l_max=4): + def __init__(self, name: str, h_lm: dict, times: np.ndarray, l_max: int = 4): self.name = name self.h_lm = h_lm @@ -22,14 +23,14 @@ def h_to_geo(self): if self.MTot is None or self.distance is None: return 1 else: - return self.distance * MPC / self.MTot / SOLAR_MASS / GG * CC ** 2 + return self.distance * MPC / self.MTot / SOLAR_MASS / GG * CC**2 @property def t_to_geo(self): if self.MTot is None or self.distance is None: return 1 else: - return 1 / self.MTot / SOLAR_MASS / GG * CC ** 3 + return 1 / self.MTot / SOLAR_MASS / GG * CC**3 @property def modes(self): @@ -47,7 +48,13 @@ def distance(self, distance): def delta_t(self): return self.times[1] - self.times[0] - def time_domain_memory(self, inc=None, phase=None, modes=None, gamma_lmlm=None): + def time_domain_memory( + self, + inc: float = None, + phase: float = None, + modes: list = None, + gamma_lmlm: dict = None, + ) -> Tuple[dict, np.ndarray]: """ Calculate the spherical harmonic decomposition of the nonlinear memory from a dictionary of spherical mode time series @@ -94,9 +101,7 @@ def time_domain_memory(self, inc=None, phase=None, modes=None, gamma_lmlm=None): pass if gamma_lmlm is not None: - warnings.warn( - f"The gamma_lmlm argument is deprecated and will be removed." - ) + warnings.warn(f"The gamma_lmlm argument is deprecated and will be removed.") # constant terms in SI units const = 1 / 4 / np.pi @@ -107,11 +112,14 @@ def time_domain_memory(self, inc=None, phase=None, modes=None, gamma_lmlm=None): modes = lmax_modes(self.l_max) for ell, delta_m in modes: - dh_mem_dt_lm[(ell, int(delta_m))] = np.sum([ - dhlm_dt_sq[(lm1, lm2)] * analytic_gamma(lm1, lm2, ell) - for lm1, lm2 in dhlm_dt_sq.keys() - if delta_m == lm1[1] - lm2[1] - ], axis=0) * np.ones(len(self.times), dtype=complex) + dh_mem_dt_lm[(ell, int(delta_m))] = np.sum( + [ + dhlm_dt_sq[(lm1, lm2)] * analytic_gamma(lm1, lm2, ell) + for lm1, lm2 in dhlm_dt_sq.keys() + if delta_m == lm1[1] - lm2[1] + ], + axis=0, + ) * np.ones(len(self.times), dtype=complex) h_mem_lm = { lm: const * cumulative_trapezoid(dh_mem_dt_lm[lm], self.times, initial=0) for lm in dh_mem_dt_lm diff --git a/gwmemory/waveforms/mwm.py b/gwmemory/waveforms/mwm.py index 8dc9a60..8997d06 100644 --- a/gwmemory/waveforms/mwm.py +++ b/gwmemory/waveforms/mwm.py @@ -1,12 +1,21 @@ +from typing import Tuple + import numpy as np -from ..utils import CC, GG, MPC, SOLAR_MASS from ..qnms import final_mass_spin, freq_damping +from ..utils import CC, GG, MPC, SOLAR_MASS from . import MemoryGenerator class MWM(MemoryGenerator): - def __init__(self, q, total_mass=60, distance=400, name="MWM", times=None): + def __init__( + self, + q: float, + total_mass: float = 60, + distance: float = 400, + name: str = "MWM", + times: np.ndarray = None, + ): super(MWM, self).__init__(name=name, h_lm=dict(), times=times) self.name = name if q > 1: @@ -21,20 +30,22 @@ def __init__(self, q, total_mass=60, distance=400, name="MWM", times=None): times = np.linspace(-900, 100, 10001) / self.t_to_geo self.times = times - def time_domain_memory(self, inc, phase, times=None, rm=3): + def time_domain_memory( + self, inc: float, phase: float, times: np.ndarray = None, rm: float = 3 + ) -> Tuple[dict, np.ndarray]: """ Calculates the plus and cross polarisations for the minimal waveform model memory waveform: eqns (5) and (9) from Favata (2009. ApJL 159) - TODO: Impement spherical harmonic decomposition? + TODO: Implement spherical harmonic decomposition? Parameters ---------- inc: float Binary inclination angle phase: float - Binary phase at coalscence + Binary phase at coalescence times: array, optional Time array on which the memory is calculated rm: float, optional @@ -54,7 +65,7 @@ def time_domain_memory(self, inc, phase, times=None, rm=3): time_geo = times * CC - mass_solar_to_geometric = SOLAR_MASS * GG / CC ** 2 + mass_solar_to_geometric = SOLAR_MASS * GG / CC**2 m1_geo = self.m1 * mass_solar_to_geometric m2_geo = self.m2 * mass_solar_to_geometric @@ -91,41 +102,41 @@ def time_domain_memory(self, inc, phase, times=None, rm=3): TT = time_geo - tm # some quantity defined after equation (7) of Favata - trr = 5 * MM * rm ** 4 / (256 * eta * MM ** 4) + trr = 5 * MM * rm**4 / (256 * eta * MM**4) # calculate the A_{ell m n} matching coefficients. Note that # I've solved a matrix equation that solved for the three coefficients # from three equations - xi = 2 * np.sqrt(2 * np.pi / 5) * eta * MM * rm ** 2 - chi = -2 * 1j * np.sqrt(MM / rm ** 3) + xi = 2 * np.sqrt(2 * np.pi / 5) * eta * MM * rm**2 + chi = -2 * 1j * np.sqrt(MM / rm**3) A220 = ( xi * ( - sigma221 * sigma222 * chi ** 2 - + sigma221 * chi ** 3 - + sigma222 * chi ** 3 - + chi ** 4 + sigma221 * sigma222 * chi**2 + + sigma221 * chi**3 + + sigma222 * chi**3 + + chi**4 ) / ((sigma220 - sigma221) * (sigma220 - sigma222)) ) A221 = ( xi * ( - sigma220 * sigma222 * chi ** 2 - + sigma220 * chi ** 3 - + sigma222 * chi ** 3 - + chi ** 4 + sigma220 * sigma222 * chi**2 + + sigma220 * chi**3 + + sigma222 * chi**3 + + chi**4 ) / ((sigma221 - sigma220) * (sigma221 - sigma222)) ) A222 = ( xi * ( - sigma220 * sigma221 * chi ** 2 - + sigma220 * chi ** 3 - + sigma221 * chi ** 3 - + chi ** 4 + sigma220 * sigma221 * chi**2 + + sigma220 * chi**3 + + sigma221 * chi**3 + + chi**4 ) / ((sigma221 - sigma222) * (sigma220 - sigma222)) ) @@ -240,7 +251,7 @@ def time_domain_memory(self, inc, phase, times=None, rm=3): cT = np.cos(inc) h_plus_coeff = ( - 0.77 * eta * MM / (384 * np.pi) * sT ** 2 * (17 + cT ** 2) / dist_geo + 0.77 * eta * MM / (384 * np.pi) * sT**2 * (17 + cT**2) / dist_geo ) h_mem = dict(plus=h_plus_coeff * h_MWM, cross=np.zeros_like(h_MWM)) diff --git a/gwmemory/waveforms/nr.py b/gwmemory/waveforms/nr.py index 2429415..735f369 100644 --- a/gwmemory/waveforms/nr.py +++ b/gwmemory/waveforms/nr.py @@ -1,4 +1,5 @@ from copy import deepcopy +from typing import Tuple import numpy as np @@ -28,12 +29,12 @@ class SXSNumericalRelativity(MemoryGenerator): def __init__( self, - name, - modes=None, - extraction="OutermostExtraction.dir", - total_mass=None, - distance=None, - times=None, + name: str, + modes: list = None, + extraction: str = "OutermostExtraction.dir", + total_mass: float = None, + distance: float = None, + times: np.ndarray = None, ): """ Initialise SXSNumericalRelativity MemoryGenerator @@ -49,16 +50,16 @@ def __init__( h5 file. total_mass: float Lab-frame total mass of the binary in solar masses. - distace: float + distance: float Luminosity distance to the binary in MPC. times: array Time array to evaluate the waveforms on, default is time array in h5 file. """ - h_lm, times = load_sxs_waveform( - name, modes=modes, extraction=extraction + h_lm, times = load_sxs_waveform(name, modes=modes, extraction=extraction) + super(SXSNumericalRelativity, self).__init__( + name=name, h_lm=h_lm, times=times, l_max=4 ) - super(SXSNumericalRelativity, self).__init__(name=name, h_lm=h_lm, times=times, l_max=4) self.MTot = total_mass self.distance = distance @@ -70,7 +71,13 @@ def __init__( if times is not None: self.set_time_array(times) - def time_domain_oscillatory(self, times=None, modes=None, inc=None, phase=None): + def time_domain_oscillatory( + self, + times: np.ndarray = None, + modes: list = None, + inc: float = None, + phase: float = None, + ) -> Tuple[dict, np.ndarray]: """ Get the mode decomposition of the numerical relativity waveform. diff --git a/gwmemory/waveforms/surrogate.py b/gwmemory/waveforms/surrogate.py index 358efb0..62e6f22 100644 --- a/gwmemory/waveforms/surrogate.py +++ b/gwmemory/waveforms/surrogate.py @@ -1,4 +1,5 @@ from copy import deepcopy +from typing import Tuple import numpy as np @@ -34,17 +35,17 @@ class Surrogate(MemoryGenerator): def __init__( self, - q, - name="nrsur7dq2", - total_mass=None, - spin_1=None, - spin_2=None, - distance=None, - l_max=4, - modes=None, - times=None, - minimum_frequency=0, - sampling_frequency=None, + q: float, + name: str = "nrsur7dq2", + total_mass: float = None, + spin_1: Tuple[float, float, float] = None, + spin_2: Tuple[float, float, float] = None, + distance: float = None, + l_max: int = 4, + modes: list = None, + times: np.ndarray = None, + minimum_frequency: float = 0, + sampling_frequency: float = None, ): """ Initialise Surrogate MemoryGenerator @@ -81,7 +82,7 @@ def __init__( print("$ python -m pip install nrsur7dq2") raise - self.sur = NRSurrogate7dq2() + self.surrogate = NRSurrogate7dq2() if q < 1: q = 1 / q @@ -106,9 +107,11 @@ def __init__( print("$ conda install -c conda-forge gwsurrogate") raise try: - self.sur = gwsurrogate.LoadSurrogate(name) + self.surrogate = gwsurrogate.LoadSurrogate(name) except ValueError: - raise ValueError(f"Surrogate model {name} not in {gwsurrogate.SURROGATE_CLASSES}") + raise ValueError( + f"Surrogate model {name} not in {gwsurrogate.SURROGATE_CLASSES}" + ) if q < 1: q = 1 / q self.q = q @@ -136,7 +139,13 @@ def __init__( if times is not None: self.set_time_array(times) - def time_domain_oscillatory(self, times=None, modes=None, inc=None, phase=None): + def time_domain_oscillatory( + self, + times: np.ndarray = None, + modes: list = None, + inc: float = None, + phase: float = None, + ) -> Tuple[dict, np.ndarray]: """ Get the mode decomposition of the surrogate waveform. @@ -172,7 +181,7 @@ def time_domain_oscillatory(self, times=None, modes=None, inc=None, phase=None): if times is None: times = np.linspace(-900, 100, 10001) times = times / self.t_to_geo - h_lm = self.sur( + h_lm = self.surrogate( self.q, self.S1, self.S2, @@ -190,7 +199,7 @@ def time_domain_oscillatory(self, times=None, modes=None, inc=None, phase=None): delta_t = 1 / self.sampling_frequency else: delta_t = None - times, h_lm, _ = self.sur( + times, h_lm, _ = self.surrogate( q=self.q, chiA0=self.S1, chiB0=self.S2, @@ -219,7 +228,7 @@ def time_domain_oscillatory(self, times=None, modes=None, inc=None, phase=None): ) ) modes = list(set(modes).union(available_modes)) - print("Using modes {}".format(" ".join(modes))) + print(f"Using modes {' '.join(modes)}") h_lm = {(ell, m): h_lm[ell, m] for ell, m in modes} diff --git a/pyproject.toml b/pyproject.toml index df13b40..86d8757 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -16,7 +16,7 @@ authors = [ ] description = "Calculate the gravitational-wave memory for arbitrary time series" readme = "README.md" -requires-python = ">=3.7" +requires-python = ">=3.8" classifiers = [ "License :: OSI Approved :: MIT License", "Operating System :: OS Independent", diff --git a/test/angles_test.py b/test/angles_test.py index 4f83f64..7962dd9 100644 --- a/test/angles_test.py +++ b/test/angles_test.py @@ -1,6 +1,5 @@ -import pytest - import numpy as np +import pytest from sxs.waveforms import WaveformModes from sxs.waveforms.memory import Dinverse @@ -10,10 +9,14 @@ @pytest.mark.parametrize("lm", [(2, 2), (2, -2), (3, 3), (3, 1)]) def test_numeric_gamma_agrees_with_analytic(lm): numeric = np.real(angles.gamma("22", f"{lm[0]}{lm[1]}")) - analytic = np.array([angles.analytic_gamma((2, 2), lm, ell) for ell in range(2, len(numeric) + 2)]) + analytic = np.array( + [angles.analytic_gamma((2, 2), lm, ell) for ell in range(2, len(numeric) + 2)] + ) assert max(abs(numeric - analytic)) < 1e-4 numeric = np.real(angles.gamma(f"{lm[0]}{lm[1]}", "22")) - analytic = np.array([angles.analytic_gamma(lm, (2, 2), ell) for ell in range(2, len(numeric) + 2)]) + analytic = np.array( + [angles.analytic_gamma(lm, (2, 2), ell) for ell in range(2, len(numeric) + 2)] + ) assert max(abs(numeric - analytic)) < 1e-4 @@ -33,7 +36,4 @@ def test_memory_factor_matches_sxs(): h_test._metadata["spin_weight"] = 0 corrected = Dinverse(h_test).ethbar.ethbar.ndarray[0, 3:] / 2 for ell in [2, 3, 4]: - assert ( - corrected[h_test.index(ell, 0)] - == angles.memory_correction(ell) - ) + assert corrected[h_test.index(ell, 0)] == angles.memory_correction(ell) diff --git a/test/harmonics_test.py b/test/harmonics_test.py index 5b7e5fd..61f6365 100644 --- a/test/harmonics_test.py +++ b/test/harmonics_test.py @@ -1,7 +1,6 @@ -import pytest - import lal import numpy as np +import pytest from gwmemory import harmonics @@ -17,4 +16,3 @@ def test_spherical_harmonics_match_lal(lm): ) assert abs(np.real(diff)) < 1e-6 assert abs(np.imag(diff)) < 1e-6 - diff --git a/test/waveform_test.py b/test/waveform_test.py index db0a91c..725bac1 100644 --- a/test/waveform_test.py +++ b/test/waveform_test.py @@ -1,13 +1,13 @@ -import h5py import os import gwsurrogate as gws +import h5py import numpy as np import pytest from gwtools import sxs_memory import gwmemory -from gwmemory import time_domain_memory, frequency_domain_memory +from gwmemory import frequency_domain_memory, time_domain_memory from gwmemory.waveforms import Approximant, Surrogate, SXSNumericalRelativity TEST_MODELS = [ @@ -57,7 +57,9 @@ def test_waveform_multiple_runs(model): assert np.allclose(mem_1[mode], mem_2[mode]) -@pytest.mark.parametrize(("model", "name"), ([Surrogate, "NRSur7dq4"], [Approximant, "IMRPhenomT"])) +@pytest.mark.parametrize( + ("model", "name"), ([Surrogate, "NRSur7dq4"], [Approximant, "IMRPhenomT"]) +) def test_minimal_arguments(model, name): """ Test calculating memory with default arguments. @@ -90,7 +92,7 @@ def test_nr_waveform(): for mode in h_osc: grp.create_dataset( f"Y_l{mode[0]}_m{mode[1]}.dat", - data=np.vstack([times, h_osc[mode].real, h_osc[mode].imag]).T + data=np.vstack([times, h_osc[mode].real, h_osc[mode].imag]).T, ) loaded = SXSNumericalRelativity("test_waveform.h5") mem1, times_1 = test.time_domain_memory() @@ -112,6 +114,10 @@ def test_memory_matches_sxs(): for ii, mode in enumerate(modes): gwmem = h_mem[mode] sxsmem = h_mem_sxs[mode] - overlap = np.vdot(gwmem, sxsmem) / np.vdot(gwmem, gwmem) ** 0.5 / np.vdot(sxsmem, sxsmem) ** 0.5 + overlap = ( + np.vdot(gwmem, sxsmem) + / np.vdot(gwmem, gwmem) ** 0.5 + / np.vdot(sxsmem, sxsmem) ** 0.5 + ) assert overlap.real > 1 - 1e-8 assert abs(overlap.imag) < 1e-5 From 6b954ecfe0d69a0e5a15c9954ac9090c1481109a Mon Sep 17 00:00:00 2001 From: Colm Talbot Date: Mon, 14 Nov 2022 15:49:20 -0500 Subject: [PATCH 2/4] MAINT: refactor MWM --- gwmemory/waveforms/mwm.py | 159 ++++++++------------------------------ 1 file changed, 33 insertions(+), 126 deletions(-) diff --git a/gwmemory/waveforms/mwm.py b/gwmemory/waveforms/mwm.py index 8997d06..85ea5e6 100644 --- a/gwmemory/waveforms/mwm.py +++ b/gwmemory/waveforms/mwm.py @@ -88,14 +88,10 @@ def time_domain_memory( # calculate the QNM frequencies and damping times # from the fits in Table VIII of Berti et al. (2006) - omega220, tau220 = freq_damping(Mf_geo, jj, nn=0) - omega221, tau221 = freq_damping(Mf_geo, jj, nn=1) - omega222, tau222 = freq_damping(Mf_geo, jj, nn=2) - print(omega222, tau221) + omega, tau = np.array([freq_damping(Mf_geo, jj, nn=nn) for nn in range(3)]).T - sigma220 = 1j * omega220 + 1 / tau220 - sigma221 = 1j * omega221 + 1 / tau221 - sigma222 = 1j * omega222 + 1 / tau222 + sigma = 1j * omega + 1 / tau + sigma_plus_sigma_star = np.add.outer(sigma, sigma.conj()) # set the matching time to be at t = 0 tm = 0 @@ -110,104 +106,26 @@ def time_domain_memory( xi = 2 * np.sqrt(2 * np.pi / 5) * eta * MM * rm**2 chi = -2 * 1j * np.sqrt(MM / rm**3) - A220 = ( + amplitude = ( xi - * ( - sigma221 * sigma222 * chi**2 - + sigma221 * chi**3 - + sigma222 * chi**3 - + chi**4 + * chi**2 + * np.array( + [ + (sigma[(ii + 1) % 3] + chi) + * (sigma[(ii + 2) % 3] + chi) + / (sigma[ii % 3] - sigma[(ii + 1) % 3]) + / (sigma[ii % 3] - sigma[(ii + 2) % 3]) + for ii in range(3) + ] ) - / ((sigma220 - sigma221) * (sigma220 - sigma222)) - ) - A221 = ( - xi - * ( - sigma220 * sigma222 * chi**2 - + sigma220 * chi**3 - + sigma222 * chi**3 - + chi**4 - ) - / ((sigma221 - sigma220) * (sigma221 - sigma222)) - ) - A222 = ( - xi - * ( - sigma220 * sigma221 * chi**2 - + sigma220 * chi**3 - + sigma221 * chi**3 - + chi**4 - ) - / ((sigma221 - sigma222) * (sigma220 - sigma222)) ) # Calculate the coefficients in the summed term of equation (9) # from Favata (2009) this is a double sum, with each variable going # from n = 0 to 2; therefore 9 terms - coeffSum00 = ( - sigma220 - * np.conj(sigma220) - * A220 - * np.conj(A220) - / (sigma220 + np.conj(sigma220)) - ) - coeffSum01 = ( - sigma220 - * np.conj(sigma221) - * A220 - * np.conj(A221) - / (sigma220 + np.conj(sigma221)) - ) - coeffSum02 = ( - sigma220 - * np.conj(sigma222) - * A220 - * np.conj(A222) - / (sigma220 + np.conj(sigma222)) - ) - - coeffSum10 = ( - sigma221 - * np.conj(sigma220) - * A221 - * np.conj(A220) - / (sigma221 + np.conj(sigma220)) - ) - coeffSum11 = ( - sigma221 - * np.conj(sigma221) - * A221 - * np.conj(A221) - / (sigma221 + np.conj(sigma221)) - ) - coeffSum12 = ( - sigma221 - * np.conj(sigma222) - * A221 - * np.conj(A222) - / (sigma221 + np.conj(sigma222)) - ) - - coeffSum20 = ( - sigma222 - * np.conj(sigma220) - * A222 - * np.conj(A220) - / (sigma222 + np.conj(sigma220)) - ) - coeffSum21 = ( - sigma222 - * np.conj(sigma221) - * A222 - * np.conj(A221) - / (sigma222 + np.conj(sigma221)) - ) - coeffSum22 = ( - sigma222 - * np.conj(sigma222) - * A222 - * np.conj(A222) - / (sigma222 + np.conj(sigma222)) + coeffs = ( + np.outer(sigma * amplitude, sigma.conj() * amplitude.conj()) + / sigma_plus_sigma_star ) # radial separation @@ -217,42 +135,31 @@ def time_domain_memory( h_MWM = np.zeros(len(TT)) # calculate strain for TT < 0. - h_MWM[TT <= 0] = 8 * np.pi * MM / rr[TT <= 0] + post_merger = TT > 0 + h_MWM[~post_merger] = 8 * np.pi * MM / rr[~post_merger] + + TT = TT[post_merger] # calculate strain for TT > 0. - term00 = coeffSum00 * (1 - np.exp(-TT[TT > 0] * (sigma220 + np.conj(sigma220)))) - term01 = coeffSum01 * (1 - np.exp(-TT[TT > 0] * (sigma220 + np.conj(sigma221)))) - term02 = coeffSum02 * (1 - np.exp(-TT[TT > 0] * (sigma220 + np.conj(sigma222)))) - - term10 = coeffSum10 * (1 - np.exp(-TT[TT > 0] * (sigma221 + np.conj(sigma220)))) - term11 = coeffSum11 * (1 - np.exp(-TT[TT > 0] * (sigma221 + np.conj(sigma221)))) - term12 = coeffSum12 * (1 - np.exp(-TT[TT > 0] * (sigma221 + np.conj(sigma222)))) - - term20 = coeffSum20 * (1 - np.exp(-TT[TT > 0] * (sigma222 + np.conj(sigma220)))) - term21 = coeffSum21 * (1 - np.exp(-TT[TT > 0] * (sigma222 + np.conj(sigma221)))) - term22 = coeffSum22 * (1 - np.exp(-TT[TT > 0] * (sigma222 + np.conj(sigma222)))) - - sum_terms = np.real( - term00 - + term01 - + term02 - + term10 - + term11 - + term12 - + term20 - + term21 - + term22 + terms = coeffs * ( + 1 - np.exp(-TT[:, np.newaxis, np.newaxis] * sigma_plus_sigma_star) ) - h_MWM[TT > 0] = 8 * np.pi * MM / rm + sum_terms / (eta * MM) + alt_sum_terms = terms.sum(axis=(1, 2)) - # calculate the plus polarisation of GWs: eqn. (5) from Favata (2009) - sT = np.sin(inc) - cT = np.cos(inc) + h_MWM[post_merger] = 8 * np.pi * MM / rm + alt_sum_terms / (eta * MM) + # calculate the plus polarisation of GWs: eqn. (5) from Favata (2009) h_plus_coeff = ( - 0.77 * eta * MM / (384 * np.pi) * sT**2 * (17 + cT**2) / dist_geo + 0.77 + * eta + * MM + / (384 * np.pi) + * np.sin(inc) ** 2 + * (17 + np.cos(inc) ** 2) + / dist_geo ) h_mem = dict(plus=h_plus_coeff * h_MWM, cross=np.zeros_like(h_MWM)) + h_mem["plus"] -= h_mem["plus"][0] return h_mem, times From 52ed48b2a5838a855d6253ab342c062bab4cc034 Mon Sep 17 00:00:00 2001 From: Colm Talbot Date: Mon, 14 Nov 2022 15:52:00 -0500 Subject: [PATCH 3/4] MAINT: add git blame ignore file --- .git-blame-ignore-revs | 0 1 file changed, 0 insertions(+), 0 deletions(-) create mode 100644 .git-blame-ignore-revs diff --git a/.git-blame-ignore-revs b/.git-blame-ignore-revs new file mode 100644 index 0000000..e69de29 From 34b86930ee36715596a60abbb7da302ed4ef8e6d Mon Sep 17 00:00:00 2001 From: Colm Talbot Date: Mon, 14 Nov 2022 16:05:00 -0500 Subject: [PATCH 4/4] MAINT: add git blame ignore file --- .git-blame-ignore-revs | 3 +++ 1 file changed, 3 insertions(+) diff --git a/.git-blame-ignore-revs b/.git-blame-ignore-revs index e69de29..98099a5 100644 --- a/.git-blame-ignore-revs +++ b/.git-blame-ignore-revs @@ -0,0 +1,3 @@ +# .git-blame-ignore-revs +# add pre-commits +73de5cd8b6c155d48ccd45e52449278e6c00657e \ No newline at end of file