# -*- coding: utf-8 -*-
"""World Magnetic Model Format tools.
Classes and function for reading and writing files in World Magnetic
Model Format and in IGRF text format.
See: https://geographiclib.sourceforge.io/html/magnetic.html#magneticformat
and https://www.ngdc.noaa.gov/IAGA/vmod/igrf.html.
"""
import os
import typing
import struct
import fnmatch
import pathlib
import datetime
import warnings
from typing import Optional, List
from collections import OrderedDict, namedtuple
from dataclasses import dataclass
from urllib.parse import urlsplit
from urllib.request import urlopen
import numpy as np
from ._typing import PathType
__all__ = ['MetaData', 'SphCoeffSet', 'WmmData', 'import_igrf_txt']
SphCoeffSet = namedtuple('SphCoeffSet', ['C', 'S'])
# @COMPATIBILITY: typing.OrderedDict is new in Python v3.7.2
if hasattr(typing, 'OrderedDict'):
SphCoeffsType = typing.OrderedDict[str, SphCoeffSet]
else:
SphCoeffsType = typing.Dict[str, SphCoeffSet]
[docs]class WmmData:
"""Magnetic field data."""
# @COMPATIBILITY: output type is a string (forward declaration)
# Python 3.7 implements PEP 563:
# "Postponed evaluation of annotations"
#
# from __future__ import annotations
def __init__(self, filename: Optional[PathType] = None) -> None:
self.metadata = MetaData()
self.coeffs = OrderedDict()
if filename:
filename = pathlib.Path(filename)
self.load(filename)
def _check(self) -> None:
years = list(self.coeffs.keys())
rate = years.pop()
if rate != 'rate':
raise RuntimeError('rate coefficients not found')
if years != [str(year) for year in self.metadata.get_years()]:
raise RuntimeError('coefficient years does not match metadata')
@staticmethod
def _load_sph_coeff_set(fd: typing.IO) -> SphCoeffSet:
data = fd.read(2 * 4)
n, m = struct.unpack('<ii', data)
nc = (m + 1) * (2 * n - m + 2) // 2
data = np.fromfile(fd, dtype=np.float64, count=nc)
C = np.zeros((n + 1, m + 1))
C.T[np.triu_indices(m + 1, 0, n + 1)] = data
nc = m * (2 * n - m + 1) // 2
data = np.fromfile(fd, dtype=np.float64, count=nc)
S = np.zeros((n + 1, m + 1))
S.T[1:, 1:][np.triu_indices(m, 0, n)] = data
return SphCoeffSet(C, S)
def _load_sph_coeffs(self, filename: PathType) -> SphCoeffsType:
"""Load spherical harmonics coeffs from a binary file in WMM format."""
filename = pathlib.Path(filename)
years = self.metadata.get_years()
with open(filename, 'rb') as fd:
id_ = fd.read(8).decode('utf-8')
if id_ != self.metadata.ID:
raise RuntimeError(
f'data ID ({id_}) does not match '
f'metadata ID ({self.metadata.ID})')
data = OrderedDict()
for year in years:
data[str(year)] = self._load_sph_coeff_set(fd)
data['rate'] = self._load_sph_coeff_set(fd)
return data
[docs] def load(self, filename: PathType) -> None:
"""Load metadata and spherical harmonics coefficients."""
filename = pathlib.Path(filename)
# self.metadata = MetaData() # @TODO: check
self.metadata.load(filename)
bin_filename = filename.with_suffix(filename.suffix + '.cof')
self.coeffs = self._load_sph_coeffs(bin_filename)
self._check()
@staticmethod
def _save_sph_coeff_set(fd: typing.IO, coeffs: SphCoeffSet) -> None:
if coeffs.C.ndim != 2:
raise TypeError('coefficients are not 2d arrays')
if coeffs.C.shape != coeffs.S.shape:
raise ValueError(
'C and S coefficient do not have the same shape '
'(C: {}, S: {})'.format(coeffs.C.shape, coeffs.S.shape))
n, m = coeffs.C.shape
n -= 1
m -= 1
if m > n:
raise TypeError(
f'invalid shape of coefficient arrays: '
f'n={n}, m={m}, expected m <= n')
# compute the effective size (n, m)
M = np.abs(coeffs.C) + np.abs(coeffs.S)
n = np.where(np.sum(M, 0) > 0)[0][-1]
m = np.where(np.sum(M, 1) > 0)[0][-1]
if m > n:
m = n
data = struct.pack('<ii', n, m)
fd.write(data)
data = coeffs.C.T[np.triu_indices(m + 1, 0, n + 1)]
data = np.ascontiguousarray(data, dtype='<f8')
fd.write(data.tobytes())
data = coeffs.S.T[1:, 1:][np.triu_indices(m, 0, n)]
data = np.ascontiguousarray(data, dtype='<f8')
fd.write(data.tobytes())
def _save_sph_coeffs(self, filename: PathType) -> None:
"""Store spherical harmonics coefficients in binary WMM format."""
id_ = self.metadata.get_id()
if len(id_) != 8:
raise ValueError(f'invalid id_ ({id_:r}): expected len(id_) == 8')
with open(filename, 'wb') as fd:
fd.write(id_.encode('utf-8'))
for coeffs in self.coeffs.values():
self._save_sph_coeff_set(fd, coeffs)
[docs] def save(self, outpath: PathType, force: bool = False) -> None:
"""Save data in WMM format (metadata and binary)."""
outpath = pathlib.Path(outpath)
if outpath.is_dir():
filename = outpath / (self.metadata.Name + '.wmm')
else:
filename = outpath
bin_filename = filename.with_suffix(filename.suffix + '.cof')
if not force:
if filename.exists():
raise RuntimeError(f'"{filename}" already exists')
if bin_filename.exists():
raise RuntimeError(f'"{bin_filename}" already exists')
self.metadata.save(filename)
self._save_sph_coeffs(bin_filename)
def _metadata_from_txt_header(header: str) -> MetaData:
"""Build a MataData object from the header of a coeffs file
in text IGRF format."""
parts = header.split()
if parts[:3] != ['g/h', 'n', 'm'] or '-' not in parts[-1]:
raise ValueError(f'invalid header line: {header!r}')
years = [float(item) for item in parts[3:-1]]
# check uniform time sampling
dyears = [b - a for a, b in zip(years[1:], years[:-1])]
if max(dyears) != min(dyears):
raise RuntimeError('non uniform time sampling detected')
metadata = MetaData()
metadata.ConversionDate = datetime.date.today().strftime(MetaData._DATEFMT)
metadata.NumModels = len(years)
metadata.Epoch = int(years[0])
metadata.DeltaEpoch = int(years[1] - years[0])
metadata.MinTime = int(years[0])
metadata.MaxTime = int(years[-1]) + metadata.DeltaEpoch
return metadata
[docs]def import_igrf_txt(path: PathType) -> WmmData:
"""Decode a spherical harmonics coefficient file in IGRF text format.
:param path:
the path to a local filename or a remote URL.
"""
urlobj = urlsplit(os.fspath(path))
filename = pathlib.Path(urlobj.path)
pattern = 'igrf[0-9][0-9]coeffs.txt'
if not fnmatch.fnmatch(filename.name, pattern):
raise ValueError(
f'invalid file name ("{filename}"), expected pattern: {pattern!r}')
if urlobj.scheme in ('', 'file'):
fd = open(filename)
else:
fd = urlopen(path)
with fd:
for line in fd:
if line.startswith('#'):
continue
elif line.startswith('g/h'):
assert filename.stem.endswith('coeffs')
metadata = _metadata_from_txt_header(line)
metadata.Name = filename.stem[:-6]
break
else:
raise RuntimeError(f'header line not found in {filename}')
years = metadata.get_years()
dtype = [
('type', 'S1'),
('n', 'int'),
('m', 'int'),
] + [(f'{year}', 'float64') for year in years] + [
('rate', 'float64')
]
dtype = np.dtype(dtype)
coeffs = np.loadtxt(fd, dtype=dtype)
n = np.max(coeffs['n'])
m = np.max(coeffs['m'])
g_idx = coeffs['type'] == b'g'
g = coeffs[g_idx]
h_idx = coeffs['type'] == b'h'
h = coeffs[h_idx]
data = OrderedDict()
for year in years:
C = np.zeros((n + 1, m + 1))
C[g['n'], g['m']] = g[str(year)]
S = np.zeros((n + 1, m + 1))
S[h['n'], h['m']] = h[str(year)]
data[str(year)] = SphCoeffSet(C, S)
C = np.zeros((n + 1, m + 1))
C[g['n'], g['m']] = g['rate']
S = np.zeros((n + 1, m + 1))
S[h['n'], h['m']] = h['rate']
data['rate'] = SphCoeffSet(C, S)
return WmmData.from_metadata_and_coeffs(metadata, data)