Merge pull request #53 from sdfordham/add-w-names-attrib

W_names allows for correct naming of weights series

pysyncon/base.py

@@ -1,5 +1,5 @@
 from __future__ import annotations
-from typing import Optional, Literal
+from typing import Optional, Literal, Sequence
 from abc import ABCMeta, abstractmethod
 
 import numpy as np
@@ -16,6 +16,7 @@ class BaseSynth(metaclass=ABCMeta):
     def __init__(self) -> None:
         self.dataprep: Optional[Dataprep] = None
         self.W: Optional[np.ndarray] = None
+        self.W_names: Optional[Sequence] = None
 
     @abstractmethod
     def fit(*args, **kwargs) -> None:
@@ -207,7 +208,7 @@ def weights(self, round: int = 3, threshold: Optional[float] = None) -> pd.Serie
         if self.W is None:
             raise ValueError("No weight matrix available; fit data first.")
         if self.dataprep is None:
-            weights_ser = pd.Series(self.W, name="weights")
+            weights_ser = pd.Series(self.W, index=self.W_names, name="weights")
         else:
             weights_ser = pd.Series(
                 self.W, index=list(self.dataprep.controls_identifier), name="weights"

pysyncon/synth.py

@@ -178,6 +178,7 @@ def fun(x):
         loss_V = self.calc_loss_V(W=W, Z0=Z0_arr, Z1=Z1_arr)
 
         self.W, self.loss_W, self.V, self.loss_V = W, loss_W, V_mat.diagonal(), loss_V
+        self.W_names = Z0.columns
 
     @staticmethod
     def calc_loss_V(W: np.ndarray, Z0: np.ndarray, Z1: np.ndarray) -> float: