add apriori parameter fixing functionality with unit test

Author: strawpants

src/xinv/core/tools.py

@@ -3,6 +3,8 @@
 
 import numpy as np
 import xarray as xr
+from xinv.core.attrs import find_xinv_unk_coord,get_state,get_type,xinv_tp,xinv_st,find_xinv_group_coords
+from xinv.core.logging import xinvlogger
 
 def find_ilocs(dsneq,dim,elements):
     idxname=f"_{dim}_idx"
@@ -27,3 +29,80 @@ def find_overlap_coords(coord1,coord2):
     uniq2=np.setdiff1d(coord2, coord1, assume_unique=False)
 
     return uniq1, intersect, uniq2
+
+
+def find_unk_idx(dsneq,sort=True,**kwargs):
+    """ 
+        Find the indices of a set of unknown parameters in the unknown vector of a normal equation system 
+        Parameters:
+        -----------
+        dsneq: xarray.Dataset
+            Dataset containing the normal equation system
+        sort: bool, optional
+            If True, the output indices are sorted in ascending order. The default is True.
+        kwargs: dict
+            keyword arguments with dimension names as keys and the elements to find as values
+            e.g. poly=[0,1],harmonics_seasonal=[1,2]
+        Returns:
+        --------
+        idxfound: np.ndarray or None
+            Indices of the found unknown parameters in the unknown vector, or None if no parameters were found
+        idxremaining: np.ndarray or None
+            
+            Indices of the unknown parameters that are complementary to the ones found
+        idxnotfound: np.ndarray or None
+            Indices of the requested unknown parameters that were not found in the system
+            
+    """
+
+    xunk_co=find_xinv_unk_coord(dsneq)
+    
+    unkdim=xunk_co.dims[0]
+    
+    group_id_co=None
+    group_seq_co=None
+
+    notfound=[]
+    found=[]
+    remaining=[]
+    for coname,searchparams in kwargs.items():
+        co_search=dsneq[coname]
+        if get_type(co_search) != xinv_tp.unk_co:
+            raise ValueError(f"Missing xinv_type: supplied coordinate name {coname} has no valid relation with unknown coordinate {xunk_co.name}")
+        dimname=co_search.dims[0]
+        if searchparams is not type(xr.DataArray):
+            #turn into DataArray
+            searchparams=xr.DataArray(searchparams,dims=dimname)
+        #find unique and overlapping coordinates over the unknown dimension
+        notfnd,fnd,remng=find_overlap_coords(searchparams,co_search)
+       
+        if get_state(co_search) == xinv_st.unlinked:
+            #we may have to apply an additional lookup in the group unknown multiindex
+            if group_id_co is None and group_seq_co is None:
+                group_id_co,group_seq_co=find_xinv_group_coords(dsneq)
+            notfound.extend([(coname,i) for i in find_ilocs(dsneq,coname,notfnd)])
+            found.extend([(coname,i) for i in find_ilocs(dsneq,coname,fnd)])
+            remaining.extend([(coname,i) for i in find_ilocs(dsneq,coname,remng)])
+
+        elif get_state(co_search) == xinv_st.linked:
+            notfound.extend(notfnd)
+            found.extend(fnd)
+            remaining.extend(remng)
+        else:
+            raise ValueError(f"Reduction coordinate {coname} has no valid link state")
+    
+    
+
+    #index vector of the found parameters
+    idxfound=find_ilocs(dsneq,unkdim,found) if len(found) > 0 else None
+    idxnotfound=find_ilocs(dsneq,unkdim,notfound) if len(notfound) > 0 else None
+    idxremaining=find_ilocs(dsneq,unkdim,remaining) if len(remaining) > 0 else None
+    if sort:
+        if idxfound is not None:
+            idxfound=np.sort(idxfound)
+        if idxremaining is not None:
+            idxremaining=np.sort(idxremaining)
+        if idxnotfound is not None:
+            idxnotfound=np.sort(idxnotfound)
+    
+    return idxfound,idxremaining,idxnotfound

src/xinv/linalg/inplace.py

@@ -9,9 +9,9 @@
 
 from scipy.linalg import cholesky
 from scipy.linalg.lapack import dpotri
-from scipy.linalg.blas import dtrsm,dsyrk
-
+from scipy.linalg.blas import dtrsm,dsyrk,dsymm
 
+import ctypes
 from enum import Enum
 
 class MemLayout(Enum):
@@ -184,3 +184,71 @@ def dsyrk_inplace(N,A,trans=0,beta=0.0,alpha=1.0):
         A[()]=adat
 
     return N
+
+def dsymm_inplace(A:xr.DataArray,B:xr.DataArray,C:xr.DataArray,alpha=1.0,beta=0.0):
+    """
+    Symmetric matrix times matrix multiplication
+        C= alpha*A*B +beta*C
+    """
+    #
+    #    C := alpha*A*B + beta*C, (side=0)
+    #or
+    #    C := alpha*B*A + beta*C, (side=1)
+    
+
+    #some basic checks
+    if A.shape[0] != A.shape[1]:
+        raise ValueError("Matrix A must be square")
+    if A.shape[1] != B.shape[0]:
+        raise ValueError("Inner dimensions of A and B do not match")
+    if A.shape[0] != C.shape[0]:
+        raise ValueError("Output matrix C rows do not match matrix A")
+    if B.shape[1] != C.shape[1]:
+        raise ValueError("Columns of C  do not match that of B")
+    
+    
+
+    restore=False
+    c_layout=memlayout(C.data)
+    
+    if c_layout == MemLayout.F_cont:
+        cdat=C.data
+        side=0 #0 means left
+    elif c_layout == MemLayout.C_cont:
+        cdat=C.data.T
+        #switch side to allow for in place operation
+        side=1
+    else:
+        xinvlogger.warning("C matrix is not C or F contiguous, applying copy and restore")
+        cdat=C.data.copy(order='F')
+        side=0 #0 means left  
+        restore=True
+
+    lower=islower(A)
+    A_layout=memlayout(A.data)
+    if A_layout == MemLayout.C_cont:
+        #To prevent an additional copy by f2py we can fake a F contigous array by transposing and switching lower/upper
+        lower=1-lower
+        adat=A.data.T
+    else:
+        adat=A.data
+    
+    
+    if side == 1:
+        bdat=B.data.T
+    else:
+        bdat=B.data
+    
+    b_layout=memlayout(B.data)
+    
+
+    if b_layout != c_layout:
+        xinvlogger.warning("B matrix contiguousness is inconsistent with C, a copy will be made by f2py")
+    #dsymm(alpha, a, b[, beta, c, side, lower, overwrite_c])
+    dsymm(alpha,adat,bdat,beta=beta,c=cdat,side=side,lower=lower,overwrite_c=1)
+
+
+    if restore:    
+        C.data[()]=cdat
+
+    return C #although changed in place

src/xinv/neq/__init__.py

@@ -3,3 +3,6 @@
 from .solve import *
 from .add import *
 from .reduce import *
+from .set_apriori import *
+from .fix import *
+

src/xinv/neq/fix.py

@@ -0,0 +1,100 @@
+## Permissions: See the  xinv  license file https://raw.githubusercontent.com/strawpants/xinv/master/LICENSE
+## Copyright (c) 2025 Roelof Rietbroek, r.rietbroek@utwente.nl
+
+
+from xinv.core.attrs import find_component,xinv_tp,find_xinv_group_coords,get_xunk_size_coname
+
+from xinv.core.tools import find_unk_idx
+import numpy as np
+from xinv.core.logging import xinvlogger
+
+
+def ifix(dsneq,idx,keep=False):
+    """
+        Fix/remove parameters in a normal equation system to their apriori values using indexing
+        Parameters:
+        -----------
+        dsneq: xarray.Dataset
+            Dataset containing the normal equation system to be fixed
+        idx: list or np.ndarray 
+            index of the coordinates to fix or keep in the system 
+
+    """
+    
+    u_sz,unkdim=get_xunk_size_coname(dsneq)
+    #compute the complementary index
+    if keep:
+        idxkeep=idx
+        if idxkeep.dtype == bool:
+            idxfix=~idxkeep
+        else:
+            idxfix=~np.isin(np.arange(u_sz),idxkeep)
+    else:
+        idxfix=idx
+        if idxfix.dtype == bool:
+            idxkeep=~idxfix
+        else:
+            idxkeep=~np.isin(np.arange(u_sz),idxfix)
+
+    o_dsneq=dsneq.sel({unkdim:idxkeep,unkdim+'_':idxkeep})
+    io_npara=find_component(o_dsneq,xinv_tp.npara)
+    #update amount of unknown parameters
+    io_npara[()]-=len(idxfix)
+
+    return o_dsneq
+
+
+def fix(dsneq, keep=False,**kwargs):
+    """
+        Fix/remove parameters from a Normal equation system using coordinate labelling
+        Parameters:
+        -----------
+        dsneq: xarray.Dataset
+            Dataset containing the normal equation system to be fixed
+        keep: bool, optional
+            If True, the parameters are kept instead of fixed. The default is False.
+        **kwargs: 
+            keyword arguments with the dimension name as key and a list of coordinate labels to be fixed/removed from the system
+            coord1 = fixlabels1 , .. coord2 = fixlabels2 
+
+    """
+    
+    idxfound,idxremaining,idxnotfound=find_unk_idx(dsneq,**kwargs)
+    if idxnotfound is not None:
+        xinvlogger.warning(f"Fix parameters contain values {idxnotfound} which are not found in the input normal equation system, ignoring those")
+
+    if (not keep and idxremaining is None) or (keep and idxfound is None):
+        xinvlogger.warning("Nothing to fix, returning input")
+        return dsneq
+    elif (not keep and idxfound is None) or (keep and idxremaining is None):
+        #cannot fix all unknowns
+        raise ValueError("Fix parameters contain all unknown parameters, cannot fix all unknowns")
+    
+    return ifix(dsneq,idx=idxfound,keep=keep)
+
+def groupfix(dsneq,groupname,keep=False):
+    """
+    Fix/remove by groupname a group of parameters from a normal equation system
+    Parameters
+    ----------
+    dsneq : xr.Dataset
+        Dataset containing the normal equation system to be fixed/removed
+    groupname : str
+        The groupname of the parameters to be fixed/removed
+    keep : bool, optional
+        If True, the group parameters are kept instead of fixed. The default is False.
+    """
+
+
+    #test whether the groupname actually exists
+    grpid_co,grpseq_co=find_xinv_group_coords(dsneq)
+    if grpid_co is None or grpseq_co is None:
+        raise ValueError("No group coordinates found in the normal equation system")
+
+    if groupname not in grpid_co.data:
+        raise ValueError(f"Groupname {groupname} not found in the normal equation system")
+
+    idx=grpid_co.data == groupname
+    return ifix(dsneq,idx=idx,keep=keep)
+
+

src/xinv/neq/neq.py

@@ -59,17 +59,6 @@ def zeros(rhsdims,coords,lower=0,norder='C'):
 
 
 
-def reduce(dsneq:xr.Dataset, idx):
-    #tbd reduce (implicitly solve) variables from a normal equation system spanned by idx
-    raise NotImplementedError("Reduce operation not yet implemented")
-
-def fix(dsneq:xr.Dataset, idx):
-    #tbd fix and (remove) solve) variables from a normal equation system spanned by idx
-    raise NotImplementedError("Fix operation not yet implemented")
-
-def set_apriori(dsneq:xr.Dataset, daapri:xr.DataArray):
-    #tbd set/change apriori values in a normal equation system
-    raise NotImplementedError("Set apriori values not yet implemented")
 
 
 

src/xinv/neq/reduce.py

@@ -4,7 +4,7 @@
 
 from xinv.core.attrs import find_xinv_unk_coord,find_neq_components,islower,xinv_st,get_state,get_type,xinv_tp,find_xinv_group_coords,get_xunk_size_coname
 
-from xinv.core.tools import find_overlap_coords,find_ilocs
+from xinv.core.tools import find_overlap_coords,find_ilocs,find_unk_idx
 import numpy as np
 import xarray as xr
 from xinv.core.logging import xinvlogger
@@ -128,59 +128,18 @@ def reduce(dsneq, keep=False,**kwargs):
 
     """
 
+    idxfound,idxremaining,idxnotfound=find_unk_idx(dsneq,**kwargs)
+    if idxnotfound is not None:
+        xinvlogger.warning(f"Reduction parameters contain values {idxnotfound} which are not found in the input normal equation system, ignoring those")
 
-    xunk_co=find_xinv_unk_coord(dsneq)
-    
-    unkdim=xunk_co.dims[0]
-    
-    group_id_co=None
-    group_seq_co=None
-
-    reduniq=[]
-    common=[]
-    neuniq=[]
-    for coname,redparams in kwargs.items():
-        co_search=dsneq[coname]
-        if get_type(co_search) != xinv_tp.unk_co:
-            raise ValueError(f"Missing xinv_type: Reduction coordinate {coname} has no valid relation with unknown coordinate {xunk_co.name}")
-        dimname=co_search.dims[0]
-        if redparams is not type(xr.DataArray):
-            #turn into DataArray
-            redparams=xr.DataArray(redparams,dims=dimname)
-        #find unique and overlapping coordinates over the unknown dimension
-        redu,com,neun=find_overlap_coords(redparams,co_search)
-       
-        if get_state(co_search) == xinv_st.unlinked:
-            #we may have to apply an additional lookup in the group unknown multiindex
-            if group_id_co is None and group_seq_co is None:
-                group_id_co,group_seq_co=find_xinv_group_coords(dsneq)
-            common.extend([(coname,i) for i in find_ilocs(dsneq,coname,com)])
-            reduniq.extend([(coname,i) for i in find_ilocs(dsneq,coname,redu)])
-            neuniq.extend([(coname,i) for i in find_ilocs(dsneq,coname,neun)])
-
-        elif get_state(co_search) == xinv_st.linked:
-            reduniq.extend(redu)
-            common.extend(com)
-            neuniq.extend(neun)
-        else:
-            raise ValueError(f"Reduction coordinate {coname} has no valid link state")
-    
-    
-    
-    if len(reduniq) > 0:
-        xinvlogger.warning(f"Reduction parameters contain values {reduniq} which are not found in the input normal equation system, ignoring those")
-
-    if len(neuniq) == 0:
-        raise ValueError("Reduction parameters contain all unknown parameters, cannot reduce all unknowns")
-    
-    if len(common) == 0:
+    if (not keep and idxremaining is None) or (keep and idxfound is None):
         xinvlogger.warning("Nothing to reduce, returning input")
         return dsneq
-
-    #index vector of the to be reduced parameters
-    idxreduce=find_ilocs(dsneq,unkdim,common)
+    elif (not keep and idxfound is None) or (keep and idxremaining is None):
+        #cannot reduce all unknowns
+        raise ValueError("Reduction parameters contain all unknown parameters, cannot reduce all unknowns")
 
-    return ireduce(dsneq,idxreduce,keep)
+    return ireduce(dsneq,idxfound,keep)
 
 def groupreduce(dsneq,groupname,keep=False):
     """