From 261e435427d6b5ea0bc430f2da5debe9e22e084d Mon Sep 17 00:00:00 2001
From: David Banas <david.banas@keysight.com>
Date: Fri, 21 Jun 2024 17:05:05 -0400
Subject: [PATCH] Fixed missing post-DFE eye and faulty bathtub extrapolation.

---
 src/pybert/models/bert.py | 18 +++++++++++-------
 src/pybert/utility.py     | 10 ++--------
 2 files changed, 13 insertions(+), 15 deletions(-)

diff --git a/src/pybert/models/bert.py b/src/pybert/models/bert.py
index 1865543..c571465 100644
--- a/src/pybert/models/bert.py
+++ b/src/pybert/models/bert.py
@@ -852,18 +852,22 @@ def update_results(self):
     xs = linspace(-ui * 1.0e12, ui * 1.0e12, width)
     height = 1000
     tiny_noise = normal(scale=1e-3, size=len(chnl_out[ignore_samps:]))  # to make channel eye easier to view.
-    y_max = 1.1 * max(abs(array(self.chnl_out)))
+    y_max = 1.1 * max(abs(array(self.chnl_out[ignore_samps:])))
     eye_chnl = calc_eye(ui, samps_per_ui, height, self.chnl_out[ignore_samps:] + tiny_noise, y_max)
-    y_max = 1.1 * max(abs(array(self.rx_in)))
+    y_max = 1.1 * max(abs(array(self.rx_in[ignore_samps:])))
     eye_tx = calc_eye(ui, samps_per_ui, height, self.rx_in[ignore_samps:], y_max)
-    y_max = 1.1 * max(abs(array(self.ctle_out)))
+    y_max = 1.1 * max(abs(array(self.ctle_out[ignore_samps:])))
     eye_ctle = calc_eye(ui, samps_per_ui, height, self.ctle_out[ignore_samps:], y_max)
+    y_max = 1.1 * max(abs(array(self.dfe_out[ignore_samps:])))
     i = 0
-    while clock_times[i] <= ignore_until:
+    len_clock_times = len(clock_times)
+    while i < len_clock_times and clock_times[i] < ignore_until:
         i += 1
-        assert i < len(clock_times), "ERROR: Insufficient coverage in 'clock_times' vector."
-    y_max = 1.1 * max(abs(array(self.dfe_out)))
-    eye_dfe = calc_eye(ui, samps_per_ui, height, self.dfe_out[ignore_samps:], y_max, clock_times[i:])
+    if i >= len(clock_times):
+        self.log("ERROR: Insufficient coverage in 'clock_times' vector.")
+        eye_dfe = calc_eye(ui, samps_per_ui, height, self.dfe_out[ignore_samps:], y_max)
+    else:
+        eye_dfe = calc_eye(ui, samps_per_ui, height, self.dfe_out[ignore_samps:], y_max, array(clock_times[i:]) - ignore_until)
     self.plotdata.set_data("eye_index", xs)
     self.plotdata.set_data("eye_chnl", eye_chnl)
     self.plotdata.set_data("eye_tx", eye_tx)
diff --git a/src/pybert/utility.py b/src/pybert/utility.py
index 14e4e14..097c9b8 100644
--- a/src/pybert/utility.py
+++ b/src/pybert/utility.py
@@ -741,7 +741,7 @@ def calc_eye(ui, samps_per_ui, height, ys, y_max, clock_times=None):  # pylint:
 
     # Generate the "heat" picture array.
     img_array = zeros([height, width])
-    if clock_times:
+    if clock_times is not None:
         for clock_time in clock_times:
             start_time = clock_time - ui
             start_ix = int(start_time / tsamp)
@@ -1542,13 +1542,7 @@ def make_bathtub(centers, jit_pdf, min_val=0, rj=0, extrap=False):  # pylint: di
     """
     half_len  = len(jit_pdf) // 2
     dt        = centers[1] - centers[0]  # Bins assumed to be uniformly spaced!
-    try:
-        jit_pdf_center_of_mass = int(mean([k * pk for (k, pk) in enumerate(jit_pdf)]))
-    except Exception as err:  # pylint: disable=broad-exception-caught
-        print(f"Error finding jitter PDF center of mass: {err}", flush=True)
-        jit_pdf_center_of_mass = half_len
-    _jit_pdf = roll(jit_pdf, half_len - jit_pdf_center_of_mass)
-    zero_locs = where(fftshift(_jit_pdf) == 0)[0]
+    zero_locs = where(fftshift(jit_pdf) == 0)[0]
     ext_first = 0
     ext_last  = len(jit_pdf)
     if (extrap and len(zero_locs)):