edsl 0.1.38.dev2__py3-none-any.whl → 0.1.38.dev4__py3-none-any.whl

edsl/scenarios/ScenarioJoin.py

@@ -0,0 +1,127 @@
+from __future__ import annotations
+from typing import Union, TYPE_CHECKING
+
+# if TYPE_CHECKING:
+from edsl.scenarios.ScenarioList import ScenarioList
+from edsl.scenarios.Scenario import Scenario
+
+
+class ScenarioJoin:
+    """Handles join operations between two ScenarioLists.
+
+    This class encapsulates all join-related logic, making it easier to maintain
+    and extend with other join types (inner, right, full) in the future.
+    """
+
+    def __init__(self, left: "ScenarioList", right: "ScenarioList"):
+        """Initialize join operation with two ScenarioLists.
+
+        Args:
+            left: The left ScenarioList
+            right: The right ScenarioList
+        """
+        self.left = left
+        self.right = right
+
+    def left_join(self, by: Union[str, list[str]]) -> ScenarioList:
+        """Perform a left join between the two ScenarioLists.
+
+        Args:
+            by: String or list of strings representing the key(s) to join on. Cannot be empty.
+
+        Returns:
+            A new ScenarioList containing the joined scenarios
+
+        Raises:
+            ValueError: If by is empty or if any join keys don't exist in both ScenarioLists
+        """
+        self._validate_join_keys(by)
+        by_keys = [by] if isinstance(by, str) else by
+
+        other_dict = self._create_lookup_dict(self.right, by_keys)
+        all_keys = self._get_all_keys()
+
+        return ScenarioList(
+            self._create_joined_scenarios(by_keys, other_dict, all_keys)
+        )
+
+    def _validate_join_keys(self, by: Union[str, list[str]]) -> None:
+        """Validate join keys exist in both ScenarioLists."""
+        if not by:
+            raise ValueError(
+                "Join keys cannot be empty. Please specify at least one key to join on."
+            )
+
+        by_keys = [by] if isinstance(by, str) else by
+        left_keys = set(next(iter(self.left)).keys()) if self.left else set()
+        right_keys = set(next(iter(self.right)).keys()) if self.right else set()
+
+        missing_left = set(by_keys) - left_keys
+        missing_right = set(by_keys) - right_keys
+        if missing_left or missing_right:
+            missing = missing_left | missing_right
+            raise ValueError(f"Join key(s) {missing} not found in both ScenarioLists")
+
+    @staticmethod
+    def _get_key_tuple(scenario: Scenario, keys: list[str]) -> tuple:
+        """Create a tuple of values for the join keys."""
+        return tuple(scenario[k] for k in keys)
+
+    def _create_lookup_dict(self, scenarios: ScenarioList, by_keys: list[str]) -> dict:
+        """Create a lookup dictionary for the right scenarios."""
+        return {
+            self._get_key_tuple(scenario, by_keys): scenario for scenario in scenarios
+        }
+
+    def _get_all_keys(self) -> set:
+        """Get all unique keys from both ScenarioLists."""
+        all_keys = set()
+        for scenario in self.left:
+            all_keys.update(scenario.keys())
+        for scenario in self.right:
+            all_keys.update(scenario.keys())
+        return all_keys
+
+    def _create_joined_scenarios(
+        self, by_keys: list[str], other_dict: dict, all_keys: set
+    ) -> list[Scenario]:
+        """Create the joined scenarios."""
+        new_scenarios = []
+
+        for scenario in self.left:
+            new_scenario = {key: None for key in all_keys}
+            new_scenario.update(scenario)
+
+            key_tuple = self._get_key_tuple(scenario, by_keys)
+            if matching_scenario := other_dict.get(key_tuple):
+                self._handle_matching_scenario(
+                    new_scenario, scenario, matching_scenario, by_keys
+                )
+
+            new_scenarios.append(Scenario(new_scenario))
+
+        return new_scenarios
+
+    def _handle_matching_scenario(
+        self,
+        new_scenario: dict,
+        left_scenario: Scenario,
+        right_scenario: Scenario,
+        by_keys: list[str],
+    ) -> None:
+        """Handle merging of matching scenarios and conflict warnings."""
+        overlapping_keys = set(left_scenario.keys()) & set(right_scenario.keys())
+
+        for key in overlapping_keys:
+            if key not in by_keys and left_scenario[key] != right_scenario[key]:
+                join_conditions = [f"{k}='{left_scenario[k]}'" for k in by_keys]
+                print(
+                    f"Warning: Conflicting values for key '{key}' where "
+                    f"{' AND '.join(join_conditions)}. "
+                    f"Keeping left value: {left_scenario[key]} "
+                    f"(discarding: {right_scenario[key]})"
+                )
+
+        # Only update with non-overlapping keys from matching scenario
+        new_keys = set(right_scenario.keys()) - set(left_scenario.keys())
+        new_scenario.update({k: right_scenario[k] for k in new_keys})

edsl/scenarios/ScenarioList.py

@@ -31,6 +31,10 @@ class ScenarioListMixin(ScenarioListPdfMixin, ScenarioListExportMixin):
 class ScenarioList(Base, UserList, ScenarioListMixin):
     """Class for creating a list of scenarios to be used in a survey."""
 
+    __documentation__ = (
+        "https://docs.expectedparrot.com/en/latest/scenarios.html#scenariolist"
+    )
+
     def __init__(self, data: Optional[list] = None, codebook: Optional[dict] = None):
         """Initialize the ScenarioList class."""
         if data is not None:
@@ -241,6 +245,9 @@ class ScenarioList(Base, UserList, ScenarioListMixin):
 
         return dict_hash(self.to_dict(sort=True, add_edsl_version=False))
 
+    def __eq__(self, other: Any) -> bool:
+        return hash(self) == hash(other)
+
     def __repr__(self):
         return f"ScenarioList({self.data})"
 
@@ -282,41 +289,49 @@ class ScenarioList(Base, UserList, ScenarioListMixin):
         random.shuffle(self.data)
         return self
 
-    def _repr_html_(self) -> str:
-        from edsl.utilities.utilities import data_to_html
-
-        data = self.to_dict()
-        _ = data.pop("edsl_version")
-        _ = data.pop("edsl_class_name")
-        for s in data["scenarios"]:
-            _ = s.pop("edsl_version")
-            _ = s.pop("edsl_class_name")
-        for scenario in data["scenarios"]:
-            for key, value in scenario.items():
-                if hasattr(value, "to_dict"):
-                    data[key] = value.to_dict()
-        return data_to_html(data)
-
-    def tally(self, field) -> dict:
-        """Return a tally of the values in the field.
-
-        Example:
-
-        >>> s = ScenarioList([Scenario({'a': 1, 'b': 1}), Scenario({'a': 1, 'b': 2})])
-        >>> s.tally('b')
-        {1: 1, 2: 1}
-        """
-        return dict(Counter([scenario[field] for scenario in self]))
-
-    def sample(self, n: int, seed="edsl") -> ScenarioList:
+    def _repr_html_(self):
+        """Return an HTML representation of the AgentList."""
+        # return (
+        #     str(self.summary(format="html")) + "<br>" + str(self.table(tablefmt="html"))
+        # )
+        footer = f"<a href={self.__documentation__}>(docs)</a>"
+        return str(self.summary(format="html")) + footer
+
+    # def _repr_html_(self) -> str:
+    # from edsl.utilities.utilities import data_to_html
+
+    # data = self.to_dict()
+    # _ = data.pop("edsl_version")
+    # _ = data.pop("edsl_class_name")
+    # for s in data["scenarios"]:
+    #     _ = s.pop("edsl_version")
+    #     _ = s.pop("edsl_class_name")
+    # for scenario in data["scenarios"]:
+    #     for key, value in scenario.items():
+    #         if hasattr(value, "to_dict"):
+    #             data[key] = value.to_dict()
+    # return data_to_html(data)
+
+    # def tally(self, field) -> dict:
+    #     """Return a tally of the values in the field.
+
+    #     Example:
+
+    #     >>> s = ScenarioList([Scenario({'a': 1, 'b': 1}), Scenario({'a': 1, 'b': 2})])
+    #     >>> s.tally('b')
+    #     {1: 1, 2: 1}
+    #     """
+    #     return dict(Counter([scenario[field] for scenario in self]))
+
+    def sample(self, n: int, seed: Optional[str] = None) -> ScenarioList:
         """Return a random sample from the ScenarioList
 
         >>> s = ScenarioList.from_list("a", [1,2,3,4,5,6])
-        >>> s.sample(3)
+        >>> s.sample(3, seed = "edsl")
         ScenarioList([Scenario({'a': 2}), Scenario({'a': 1}), Scenario({'a': 3})])
         """
-
-        random.seed(seed)
+        if seed:
+            random.seed(seed)
 
         return ScenarioList(random.sample(self.data, n))
 
@@ -564,6 +579,47 @@ class ScenarioList(Base, UserList, ScenarioListMixin):
             func = lambda x: x
         return cls([Scenario({name: func(value)}) for value in values])
 
+    def table(self, *fields, tablefmt=None, pretty_labels=None) -> str:
+        """Return the ScenarioList as a table."""
+
+        from tabulate import tabulate_formats
+
+        if tablefmt is not None and tablefmt not in tabulate_formats:
+            raise ValueError(
+                f"Invalid table format: {tablefmt}",
+                f"Valid formats are: {tabulate_formats}",
+            )
+        return self.to_dataset().table(
+            *fields, tablefmt=tablefmt, pretty_labels=pretty_labels
+        )
+
+    def tree(self, node_list: Optional[List[str]] = None) -> str:
+        """Return the ScenarioList as a tree."""
+        return self.to_dataset().tree(node_list)
+
+    def _summary(self):
+        d = {
+            "EDSL Class name": "ScenarioList",
+            "# Scenarios": len(self),
+            "Scenario Keys": list(self.parameters),
+        }
+        return d
+
+    def reorder_keys(self, new_order):
+        """Reorder the keys in the scenarios.
+
+        Example:
+
+        >>> s = ScenarioList([Scenario({'a': 1, 'b': 2}), Scenario({'a': 3, 'b': 4})])
+        >>> s.reorder_keys(['b', 'a'])
+        ScenarioList([Scenario({'b': 2, 'a': 1}), Scenario({'b': 4, 'a': 3})])
+        """
+        new_scenarios = []
+        for scenario in self:
+            new_scenario = Scenario({key: scenario[key] for key in new_order})
+            new_scenarios.append(new_scenario)
+        return ScenarioList(new_scenarios)
+
     def to_dataset(self) -> "Dataset":
         """
         >>> s = ScenarioList.from_list("a", [1,2,3])
@@ -579,16 +635,32 @@ class ScenarioList(Base, UserList, ScenarioListMixin):
         data = [{key: [scenario[key] for scenario in self.data]} for key in keys]
         return Dataset(data)
 
-    def split(
-        self, field: str, split_on: str, index: int, new_name: Optional[str] = None
+    def unpack(
+        self, field: str, new_names: Optional[List[str]] = None, keep_original=True
     ) -> ScenarioList:
-        """Split a scenario fiel in multiple fields."""
-        if new_name is None:
-            new_name = field + "_split_" + str(index)
+        """Unpack a field into multiple fields.
+
+        Example:
+
+        >>> s = ScenarioList([Scenario({'a': 1, 'b': [2, True]}), Scenario({'a': 3, 'b': [3, False]})])
+        >>> s.unpack('b')
+        ScenarioList([Scenario({'a': 1, 'b': [2, True], 'b_0': 2, 'b_1': True}), Scenario({'a': 3, 'b': [3, False], 'b_0': 3, 'b_1': False})])
+        >>> s.unpack('b', new_names=['c', 'd'], keep_original=False)
+        ScenarioList([Scenario({'a': 1, 'c': 2, 'd': True}), Scenario({'a': 3, 'c': 3, 'd': False})])
+
+        """
+        new_names = new_names or [f"{field}_{i}" for i in range(len(self[0][field]))]
         new_scenarios = []
         for scenario in self:
             new_scenario = scenario.copy()
-            new_scenario[new_name] = scenario[field].split(split_on)[index]
+            if len(new_names) == 1:
+                new_scenario[new_names[0]] = scenario[field]
+            else:
+                for i, new_name in enumerate(new_names):
+                    new_scenario[new_name] = scenario[field][i]
+
+            if not keep_original:
+                del new_scenario[field]
             new_scenarios.append(new_scenario)
         return ScenarioList(new_scenarios)
 
@@ -901,33 +973,32 @@ class ScenarioList(Base, UserList, ScenarioListMixin):
         return cls.from_excel(temp_filename, sheet_name=sheet_name)
 
     @classmethod
-    def from_csv(cls, source: Union[str, urllib.parse.ParseResult]) -> ScenarioList:
-        """Create a ScenarioList from a CSV file or URL.
+    def from_delimited_file(
+        cls, source: Union[str, urllib.parse.ParseResult], delimiter: str = ","
+    ) -> ScenarioList:
+        """Create a ScenarioList from a delimited file (CSV/TSV) or URL.
 
         Args:
-            source: A string representing either a local file path or a URL to a CSV file,
+            source: A string representing either a local file path or a URL to a delimited file,
                     or a urllib.parse.ParseResult object for a URL.
+            delimiter: The delimiter used in the file. Defaults to ',' for CSV files.
+                    Use '\t' for TSV files.
 
         Returns:
-            ScenarioList: A ScenarioList object containing the data from the CSV.
+            ScenarioList: A ScenarioList object containing the data from the file.
 
         Example:
+            # For CSV files
 
-        >>> import tempfile
-        >>> import os
-        >>> with tempfile.NamedTemporaryFile(delete=False, mode='w', suffix='.csv') as f:
-        ...     _ = f.write("name,age,location\\nAlice,30,New York\\nBob,25,Los Angeles\\n")
-        ...     temp_filename = f.name
-        >>> scenario_list = ScenarioList.from_csv(temp_filename)
-        >>> len(scenario_list)
-        2
-        >>> scenario_list[0]['name']
-        'Alice'
-        >>> scenario_list[1]['age']
-        '25'
+            >>> with open('data.csv', 'w') as f:
+            ...     _ = f.write('name,age\\nAlice,30\\nBob,25\\n')
+            >>> scenario_list = ScenarioList.from_delimited_file('data.csv')
+
+            # For TSV files
+            >>> with open('data.tsv', 'w') as f:
+            ...     _ = f.write('name\\tage\\nAlice\t30\\nBob\t25\\n')
+            >>> scenario_list = ScenarioList.from_delimited_file('data.tsv', delimiter='\\t')
 
-        >>> url = "https://example.com/data.csv"
-        >>> ## scenario_list_from_url = ScenarioList.from_csv(url)
         """
         from edsl.scenarios.Scenario import Scenario
 
@@ -940,24 +1011,111 @@ class ScenarioList(Base, UserList, ScenarioListMixin):
 
         if isinstance(source, str) and is_url(source):
             with urllib.request.urlopen(source) as response:
-                csv_content = response.read().decode("utf-8")
-            csv_file = StringIO(csv_content)
+                file_content = response.read().decode("utf-8")
+            file_obj = StringIO(file_content)
         elif isinstance(source, urllib.parse.ParseResult):
             with urllib.request.urlopen(source.geturl()) as response:
-                csv_content = response.read().decode("utf-8")
-            csv_file = StringIO(csv_content)
+                file_content = response.read().decode("utf-8")
+            file_obj = StringIO(file_content)
         else:
-            csv_file = open(source, "r")
+            file_obj = open(source, "r")
 
         try:
-            reader = csv.reader(csv_file)
+            reader = csv.reader(file_obj, delimiter=delimiter)
             header = next(reader)
             observations = [Scenario(dict(zip(header, row))) for row in reader]
         finally:
-            csv_file.close()
+            file_obj.close()
 
         return cls(observations)
 
+    # Convenience methods for specific file types
+    @classmethod
+    def from_csv(cls, source: Union[str, urllib.parse.ParseResult]) -> ScenarioList:
+        """Create a ScenarioList from a CSV file or URL."""
+        return cls.from_delimited_file(source, delimiter=",")
+
+    def left_join(self, other: ScenarioList, by: Union[str, list[str]]) -> ScenarioList:
+        """Perform a left join with another ScenarioList, following SQL join semantics.
+
+        Args:
+            other: The ScenarioList to join with
+            by: String or list of strings representing the key(s) to join on. Cannot be empty.
+
+        >>> s1 = ScenarioList([Scenario({'name': 'Alice', 'age': 30}), Scenario({'name': 'Bob', 'age': 25})])
+        >>> s2 = ScenarioList([Scenario({'name': 'Alice', 'location': 'New York'}), Scenario({'name': 'Charlie', 'location': 'Los Angeles'})])
+        >>> s3 = s1.left_join(s2, 'name')
+        >>> s3 == ScenarioList([Scenario({'age': 30, 'location': 'New York', 'name': 'Alice'}), Scenario({'age': 25, 'location': None, 'name': 'Bob'})])
+        True
+        """
+        from edsl.scenarios.ScenarioJoin import ScenarioJoin
+
+        sj = ScenarioJoin(self, other)
+        return sj.left_join(by)
+        # # Validate join keys
+        # if not by:
+        #     raise ValueError(
+        #         "Join keys cannot be empty. Please specify at least one key to join on."
+        #     )
+
+        # # Convert single string to list for consistent handling
+        # by_keys = [by] if isinstance(by, str) else by
+
+        # # Verify all join keys exist in both ScenarioLists
+        # left_keys = set(next(iter(self)).keys()) if self else set()
+        # right_keys = set(next(iter(other)).keys()) if other else set()
+
+        # missing_left = set(by_keys) - left_keys
+        # missing_right = set(by_keys) - right_keys
+        # if missing_left or missing_right:
+        #     missing = missing_left | missing_right
+        #     raise ValueError(f"Join key(s) {missing} not found in both ScenarioLists")
+
+        # # Create lookup dictionary from the other ScenarioList
+        # def get_key_tuple(scenario: Scenario, keys: list[str]) -> tuple:
+        #     return tuple(scenario[k] for k in keys)
+
+        # other_dict = {get_key_tuple(scenario, by_keys): scenario for scenario in other}
+
+        # # Collect all possible keys (like SQL combining all columns)
+        # all_keys = set()
+        # for scenario in self:
+        #     all_keys.update(scenario.keys())
+        # for scenario in other:
+        #     all_keys.update(scenario.keys())
+
+        # new_scenarios = []
+        # for scenario in self:
+        #     new_scenario = {
+        #         key: None for key in all_keys
+        #     }  # Start with nulls (like SQL)
+        #     new_scenario.update(scenario)  # Add all left values
+
+        #     key_tuple = get_key_tuple(scenario, by_keys)
+        #     if matching_scenario := other_dict.get(key_tuple):
+        #         # Check for overlapping keys with different values
+        #         overlapping_keys = set(scenario.keys()) & set(matching_scenario.keys())
+        #         for key in overlapping_keys:
+        #             if key not in by_keys and scenario[key] != matching_scenario[key]:
+        #                 join_conditions = [f"{k}='{scenario[k]}'" for k in by_keys]
+        #                 print(
+        #                     f"Warning: Conflicting values for key '{key}' where {' AND '.join(join_conditions)}. "
+        #                     f"Keeping left value: {scenario[key]} (discarding: {matching_scenario[key]})"
+        #                 )
+
+        #         # Only update with non-overlapping keys from matching scenario
+        #         new_keys = set(matching_scenario.keys()) - set(scenario.keys())
+        #         new_scenario.update({k: matching_scenario[k] for k in new_keys})
+
+        #     new_scenarios.append(Scenario(new_scenario))
+
+        # return ScenarioList(new_scenarios)
+
+    @classmethod
+    def from_tsv(cls, source: Union[str, urllib.parse.ParseResult]) -> ScenarioList:
+        """Create a ScenarioList from a TSV file or URL."""
+        return cls.from_delimited_file(source, delimiter="\t")
+
     def to_dict(self, sort=False, add_edsl_version=True) -> dict:
         """
         >>> s = ScenarioList([Scenario({'food': 'wood chips'}), Scenario({'food': 'wood-fired pizza'})])
@@ -1074,8 +1232,25 @@ class ScenarioList(Base, UserList, ScenarioListMixin):
         """
         from edsl.agents.AgentList import AgentList
         from edsl.agents.Agent import Agent
+        import warnings
+
+        agents = []
+        for scenario in self:
+            new_scenario = scenario.copy().data
+            if "name" in new_scenario:
+                name = new_scenario.pop("name")
+                proposed_agent_name = "agent_name"
+                while proposed_agent_name not in new_scenario:
+                    proposed_agent_name += "_"
+                warnings.warn(
+                    f"The 'name' field is reserved for the agent's name---putting this value in {proposed_agent_name}"
+                )
+                new_scenario[proposed_agent_name] = name
+                agents.append(Agent(traits=new_scenario, name=name))
+            else:
+                agents.append(Agent(traits=new_scenario))
 
-        return AgentList([Agent(traits=s.data) for s in self])
+        return AgentList(agents)
 
     def chunk(
         self,

edsl/surveys/Survey.py

@@ -41,6 +41,8 @@ class ValidatedString(str):
 class Survey(SurveyExportMixin, SurveyFlowVisualizationMixin, Base):
     """A collection of questions that supports skip logic."""
 
+    __documentation__ = """https://docs.expectedparrot.com/en/latest/surveys.html"""
+
     questions = QuestionsDescriptor()
     """
     A collection of questions that supports skip logic.
@@ -1587,10 +1589,22 @@ class Survey(SurveyExportMixin, SurveyFlowVisualizationMixin, Base):
         # question_names_string = ", ".join([repr(name) for name in self.question_names])
         return f"Survey(questions=[{questions_string}], memory_plan={self.memory_plan}, rule_collection={self.rule_collection}, question_groups={self.question_groups})"
 
+    def _summary(self) -> dict:
+        return {
+            "EDSL Class": "Survey",
+            "Number of Questions": len(self),
+            "Question Names": self.question_names,
+        }
+
     def _repr_html_(self) -> str:
-        from edsl.utilities.utilities import data_to_html
+        footer = f"<a href={self.__documentation__}>(docs)</a>"
+        return str(self.summary(format="html")) + footer
+
+    def tree(self, node_list: Optional[List[str]] = None):
+        return self.to_scenario_list().tree(node_list=node_list)
 
-        return data_to_html(self.to_dict())
+    def table(self, *fields, tablefmt=None) -> Table:
+        return self.to_scenario_list().to_dataset().table(*fields, tablefmt=tablefmt)
 
     def rich_print(self) -> Table:
         """Print the survey in a rich format.

edsl/surveys/SurveyFlowVisualizationMixin.py

@@ -1,27 +1,85 @@
-"""A mixin for visualizing the flow of a survey."""
+"""A mixin for visualizing the flow of a survey with parameter nodes."""
 
 from typing import Optional
 from edsl.surveys.base import RulePriority, EndOfSurvey
 import tempfile
+import os
 
 
 class SurveyFlowVisualizationMixin:
-    """A mixin for visualizing the flow of a survey."""
+    """A mixin for visualizing the flow of a survey with parameter visualization."""
 
     def show_flow(self, filename: Optional[str] = None):
-        """Create an image showing the flow of users through the survey."""
+        """Create an image showing the flow of users through the survey and question parameters."""
         # Create a graph object
         import pydot
 
         graph = pydot.Dot(graph_type="digraph")
 
-        # Add nodes for each question
+        # First collect all unique parameters and answer references
+        params_and_refs = set()
+        param_to_questions = {}  # Keep track of which questions use each parameter
+        answer_refs = set()  # Track answer references between questions
+
+        # First pass: collect parameters and their question associations
         for index, question in enumerate(self.questions):
-            graph.add_node(
-                pydot.Node(
-                    f"Q{index}", label=f"{question.question_name}", shape="ellipse"
+            # Add the main question node
+            question_node = pydot.Node(
+                f"Q{index}", label=f"{question.question_name}", shape="ellipse"
+            )
+            graph.add_node(question_node)
+
+            if hasattr(question, "parameters"):
+                for param in question.parameters:
+                    # Check if this is an answer reference (contains '.answer')
+                    if ".answer" in param:
+                        answer_refs.add((param.split(".")[0], index))
+                    else:
+                        params_and_refs.add(param)
+                        if param not in param_to_questions:
+                            param_to_questions[param] = []
+                        param_to_questions[param].append(index)
+
+        # Create parameter nodes and connect them to questions
+        for param in params_and_refs:
+            param_node_name = f"param_{param}"
+            param_node = pydot.Node(
+                param_node_name,
+                label=f"{{{{ {param} }}}}",
+                shape="box",
+                style="filled",
+                fillcolor="lightgrey",
+                fontsize="10",
+            )
+            graph.add_node(param_node)
+
+            # Connect this parameter to all questions that use it
+            for q_index in param_to_questions[param]:
+                param_edge = pydot.Edge(
+                    param_node_name,
+                    f"Q{q_index}",
+                    style="dotted",
+                    color="grey",
+                    arrowsize="0.5",
                 )
+                graph.add_edge(param_edge)
+
+        # Add edges for answer references
+        for source_q_name, target_q_index in answer_refs:
+            # Find the source question index by name
+            source_q_index = next(
+                i
+                for i, q in enumerate(self.questions)
+                if q.question_name == source_q_name
+            )
+            ref_edge = pydot.Edge(
+                f"Q{source_q_index}",
+                f"Q{target_q_index}",
+                style="dashed",
+                color="purple",
+                label="answer reference",
             )
+            graph.add_edge(ref_edge)
 
         # Add an "EndOfSurvey" node
         graph.add_node(
@@ -30,7 +88,7 @@ class SurveyFlowVisualizationMixin:
 
         # Add edges for normal flow through the survey
         num_questions = len(self.questions)
-        for index in range(num_questions - 1):  # From Q1 to Q3
+        for index in range(num_questions - 1):
             graph.add_edge(pydot.Edge(f"Q{index}", f"Q{index+1}"))
 
         graph.add_edge(pydot.Edge(f"Q{num_questions-1}", "EndOfSurvey"))
@@ -64,7 +122,7 @@ class SurveyFlowVisualizationMixin:
                 if rule.next_q != EndOfSurvey and rule.next_q < num_questions
                 else "EndOfSurvey"
             )
-            if rule.before_rule:  # Assume skip rules have an attribute `is_skip`
+            if rule.before_rule:
                 edge = pydot.Edge(
                     source_node,
                     target_node,