Manipulate seismometer response and PAZ (poles and zeros) file

respaz.py

### respaz.py
### Manipulate seismometer response or PAZ (poles and zeros) file
###
### July 09, 2021  Yasu Sawaki
### July 11, 2021  `load_resp` can retrieve the total seosor sensitivity


from pathlib import Path


def load_resp(
    fname, 
    sensitivity=1.0,
    sensitivity_magnification=0,
):
    '''Load a response file to retrive PAZ (poles and zeros) dict.
    
    Parameters:
    -----------
    fname: str, pathlib.Path
        Path to the response file to load
    sensitivity: float or bool, default to 1.0
        Overall sensitivity of recording instrument to correct.
        If set to True, the total sensor sensitivity will be retrieved.
    sensitivity_magnification: int, default to 0
        Magnify the sensitivity by the order of `sensitivity_magnification`.
        If set to -9, then convert "meter" to "nanometer". 
        Valid only when `sensitivity=True`.
        
    Returns:
    --------
    paz: dict
        PAZ as dict. 'poles' and 'zeros' are complex numbers.
        'gain' is A0 normalization factor.
        'sensitivity' is the total sensor sensitivity.
    '''
    
    ## if retrieve sensitivity
    if_sensitivity = False
    if isinstance(sensitivity, bool):
        if sensitivity:
            if_sensitivity = True
        else:
            sensitivity = 1.0      
    else:    
        sensitivity = float(sensitivity)
        if sensitivity <= 0.:
            raise ValueError("'sensitivity' must be positive.")
    
    fname = Path(fname)
    if not fname.exists():
        raise FileNotFoundError(f'{fname} does not exist.')
    
    with fname.open('r') as f:
        
        poles = []
        zeros = []
        
        for fl in f.readlines():

            ## Complex zeros
            if fl[:10] == "B053F10-13":
                zeros.append(_list2complex(fl.split()[2:4]))
            ## Complex poles
            elif fl[:10] == "B053F15-18":
                poles.append(_list2complex(fl.split()[2:4]))
            ## Gain - A0 - CONSTANT
            elif fl[:7] == "B053F07":
                gain = float(fl.split()[-1])
            elif fl[:23] == "B058F04     Sensitivity":
                if if_sensitivity:
                    ## the final sensitivity
                    sensitivity = float(fl.split()[-1])
    
    if if_sensitivity:
        if (
            isinstance(sensitivity_magnification, int) 
            and (-12 <= sensitivity_magnification <= 3)
        ):
            sensitivity *= 10 ** sensitivity_magnification
        else:
            raise ValueError(
                "'sensitivity_magnification' should be a power of 10.\n" + 
                "-12 <= 'sensitivity_magnification' <= 3"
            )
    
    paz = dict(
        poles=poles,
        zeros=zeros, 
        gain=gain,
        sensitivity=sensitivity
    )
    
    return paz


def load_paz(fname, sensitivity=1.0):
    '''Load a SAC-format PAZ (poles and zeros) file to retrive PAZ dict.
    
    Parameters:
    -----------
    fname: str, pathlib.Path
        Path to the PAZ file to load
    sensitivity: float, default to 1.0
        Overall sensitivity of recording instrument to correct.
        
    Returns:
    --------
    paz: dict
        PAZ as dict. 'poles' and 'zeros' are complex numbers.
        'gain' is A0 normalization factor.
    '''
    
    sensitivity = float(sensitivity)
    if sensitivity <= 0.:
        raise ValueError("'sensitivity' must be positive.")
    
    fname = Path(fname)
    if not fname.exists():
        raise FileNotFoundError(f'{fname} does not exist.')
    
    with fname.open('r') as f:
        
        flag_poles = 0
        flag_zeros = 0
        poles = []
        zeros = []
        
        for fl in f.readlines():
            
            if flag_zeros:
                try:
                    zeros.append(_list2complex(fl.split()))
                    flag_zeros -= 1
                except ValueError:
                    zeros = [0j] * flag_zeros
                    flag_zeros = 0
                
            if flag_poles:
                poles.append(_list2complex(fl.split()))
                flag_poles -= 1   
            
            ## ZEROS
            if fl[:5].upper() == 'ZEROS':
                flag_zeros = int(fl.split()[-1])
            ## POLES
            elif fl[:5].upper() == 'POLES':
                flag_poles = int(fl.split()[-1])
            # CONSTANT
            elif fl[:5].upper() == 'CONST':
                gain = float(fl.split()[-1])
    
    paz = dict(
        poles=poles,
        zeros=zeros, 
        gain=gain,
        sensitivity=sensitivity
    )
    
    return paz


def write_paz(fname, paz):
    '''Write a PAZ (poles and zeros) file as SAC format.
    
    Parameters:
    -----------
    fname: str, pathlib.Path
        Path to the PAZ filename to write.
    paz: dict
        Input PAZ dict.
        
    '''
    
    fname = Path(fname)
    
    with fname.open('w') as f:
        
        ## ZEROS
        f.write(f'ZEROS {len(paz["zeros"])}\n')
        for z in paz["zeros"]:
            f.write(f'{z.real} {z.imag}\n')
        
        ## POLES
        f.write(f'POLES {len(paz["poles"])}\n')
        for p in paz["poles"]:
            f.write(f'{p.real} {p.imag}\n')
            
        ## CONSTANT
        f.write(f'CONSTANT {paz["gain"]*paz["sensitivity"]:.6e}')
        

def _list2complex(data):
    return complex(float(data[0]), float(data[1]))


__author__ = 'Yasu Sawaki'
__status__ = "production"
__date__ = "July 09, 2021"