Coverage for tcvx21/units_m.py: 100%

1""" 

2Set up the "Pint" unit library 

3 

4See unit_registry.txt for defined units (for reference only) 

5""" 

6import pint 

7import os 

8import warnings 

9import xarray as xr 

10import numpy as np 

11 

12# Disable Pint's old fallback behavior (must come before importing Pint) 

13# Ignore pylint warning C0413: Import "import pint" should be placed at 

14# the top of the module 

15os.environ['PINT_ARRAY_PROTOCOL_FALLBACK'] = "0" 

16 

17unit_registry = pint.UnitRegistry() 

18Quantity = unit_registry.Quantity 

19Dimensionless = Quantity(1, "") 

20 

21# Silence NEP 18 warning 

22with warnings.catch_warnings(): 

23 warnings.simplefilter("ignore") 

24 Quantity([]) 

25 

26unit_registry.setup_matplotlib() 

27 

28# Hide a warning that units are stripped when downcasting to numpy arrays 

29warnings.simplefilter('ignore', category=pint.errors.UnitStrippedWarning) 

30 

31 

32def convert_xarray_to_quantity(input_array: xr.DataArray) -> Quantity: 

33 """ 

34 Converts an xarray with units attribute into a pint.Quantity with units (useful for unit checking) 

35 

36 If the base array is already a Quantity, its units will be silently overwritten 

37 

38 Note that all attributes except units will be lost. 

39 Not recommended except as a sanity check/in tests 

40 """ 

41 with warnings.catch_warnings(): 

42 warnings.simplefilter('ignore', category=pint.errors.UnitStrippedWarning) 

43 return input_array.norm * np.asarray(input_array)