AgentSet

Container for agents into heirarchical sets (e.g. all_agents > hiv_agents)

Source code in titan/agent.py

class AgentSet:
    """
    Container for agents into heirarchical sets (e.g. all_agents > hiv_agents)
    """

    def __init__(
        self,
        id: str,
        parent: Optional["AgentSet"] = None,
    ):
        """
        Constructor of an AgentSet

        args:
            id: name of the set
            parent: the set this set is a subset of
        """
        # members stores agent set members in a dictionary keyed by ID
        self.id = id
        self.members: Set[Agent] = set()
        self.subset: Dict[str, AgentSet] = {}

        # parent_set stores the parent set if this set is a member of an
        # AgentSet class instance. For example, for a set that is a
        # member of a larger set, the _parent_set for that set  would
        # be that larger set.
        self.parent_set = parent
        if parent:
            parent.add_subset(self)

    def __repr__(self):
        return self.id

    def __str__(self):
        return self.id

    def clear_set(self):
        """
        Clears a set of any members and subsets
        """
        self.members: Set[Agent] = set()
        self.subset: Dict[str, str] = {}

    def __iter__(self) -> Iterator[Agent]:
        """
        Iterate over the memers in the set

        example:
            ```py
            for agent in agent_set:
                # do something with agent
            ```

        returns:
            iterator over member agents
        """
        return self.members.__iter__()

    def __contains__(self, item) -> bool:
        """
        Is an agent a member of this agent set

        example:
            ```py
            if agent in agent_set:
                # do something
            ```

        returns:
            whether agent is part of set
        """
        return self.members.__contains__(item)

    # adding trickles up
    def add_agent(self, agent: Agent) -> None:
        """
        Adds an agent to the set and any parent sets

        args:
            agent: agent to add
        """
        self.members.add(agent)

        if self.parent_set is not None:
            self.parent_set.add_agent(agent)

    # removing trickles down
    def remove_agent(self, agent: Agent) -> None:
        """
        Removes agent from agent set if they are a member of the set.  Also removes the agent from any subsets.

        args:
            agent: agent to remove
        """
        if agent in self.members:
            self.members.remove(agent)

        for subset in self.iter_subset():
            subset.remove_agent(agent)

    def num_members(self) -> int:
        """
        Number of members in the set

        returns:
            number of members
        """
        return len(self.members)

    def add_subset(self, subset: "AgentSet") -> None:
        """
        Adds a new AgentSet to the current sets subset.

        args:
            subset: subset to add to this set
        """
        if subset.id not in self.subset:
            self.subset[subset.id] = subset

    def iter_subset(self) -> Iterator["AgentSet"]:
        """
        Iterate over the subsets of this agent set

        returns:
            iterator of agent sets
        """
        for subset in list(self.subset.values()):
            yield subset

    def print_subsets(self, printer=print):
        """
        Pretty print the subsets of this agent set
        """
        lines = []

        lines.append(f"\t__________ {self.id} __________")
        # lines.append("\tID\t\tN\t\t%")
        lines.append("\t{:^6}\t\t{:^5}\t\t{:^4}".format("ID", "N", "%"))
        for set in self.iter_subset():
            lines.append(
                "\t{:^6}\t\t{:^5}\t\t{:.2}".format(
                    set.id,
                    set.num_members(),
                    safe_divide(set.num_members(), set.parent_set.num_members()),
                )
            )
            for subset in set.iter_subset():
                lines.append(
                    "\t{:4}\t\t{:5}\t\t{:.2}".format(
                        subset.id,
                        subset.num_members(),
                        safe_divide(
                            subset.num_members(), subset.parent_set.num_members()
                        ),
                    )
                )
        lines.append("\t______________ END ______________")
        printer("\n".join(lines))

__contains__(item)

Is an agent a member of this agent set

example

if agent in agent_set:
    # do something

Returns:

Type	Description
bool	whether agent is part of set

Source code in titan/agent.py

def __contains__(self, item) -> bool:
    """
    Is an agent a member of this agent set

    example:
        ```py
        if agent in agent_set:
            # do something
        ```

    returns:
        whether agent is part of set
    """
    return self.members.__contains__(item)

__init__(id, parent=None)

Constructor of an AgentSet

Parameters:

Name	Type	Description	Default
id	str	name of the set	required
parent	Optional[AgentSet]	the set this set is a subset of	None

Source code in titan/agent.py

def __init__(
    self,
    id: str,
    parent: Optional["AgentSet"] = None,
):
    """
    Constructor of an AgentSet

    args:
        id: name of the set
        parent: the set this set is a subset of
    """
    # members stores agent set members in a dictionary keyed by ID
    self.id = id
    self.members: Set[Agent] = set()
    self.subset: Dict[str, AgentSet] = {}

    # parent_set stores the parent set if this set is a member of an
    # AgentSet class instance. For example, for a set that is a
    # member of a larger set, the _parent_set for that set  would
    # be that larger set.
    self.parent_set = parent
    if parent:
        parent.add_subset(self)

__iter__()

Iterate over the memers in the set

example

for agent in agent_set:
    # do something with agent

Returns:

Type	Description
Iterator[Agent]	iterator over member agents

Source code in titan/agent.py

def __iter__(self) -> Iterator[Agent]:
    """
    Iterate over the memers in the set

    example:
        ```py
        for agent in agent_set:
            # do something with agent
        ```

    returns:
        iterator over member agents
    """
    return self.members.__iter__()

add_agent(agent)

Adds an agent to the set and any parent sets

Parameters:

Name	Type	Description	Default
agent	Agent	agent to add	required

Source code in titan/agent.py

def add_agent(self, agent: Agent) -> None:
    """
    Adds an agent to the set and any parent sets

    args:
        agent: agent to add
    """
    self.members.add(agent)

    if self.parent_set is not None:
        self.parent_set.add_agent(agent)

add_subset(subset)

Adds a new AgentSet to the current sets subset.

Parameters:

Name	Type	Description	Default
subset	AgentSet	subset to add to this set	required

Source code in titan/agent.py

def add_subset(self, subset: "AgentSet") -> None:
    """
    Adds a new AgentSet to the current sets subset.

    args:
        subset: subset to add to this set
    """
    if subset.id not in self.subset:
        self.subset[subset.id] = subset

clear_set()

Clears a set of any members and subsets

Source code in titan/agent.py

def clear_set(self):
    """
    Clears a set of any members and subsets
    """
    self.members: Set[Agent] = set()
    self.subset: Dict[str, str] = {}

iter_subset()

Iterate over the subsets of this agent set

Returns:

Type	Description
Iterator[AgentSet]	iterator of agent sets

Source code in titan/agent.py

def iter_subset(self) -> Iterator["AgentSet"]:
    """
    Iterate over the subsets of this agent set

    returns:
        iterator of agent sets
    """
    for subset in list(self.subset.values()):
        yield subset

num_members()

Number of members in the set

Returns:

Type	Description
int	number of members

Source code in titan/agent.py

def num_members(self) -> int:
    """
    Number of members in the set

    returns:
        number of members
    """
    return len(self.members)

print_subsets(printer=print)

Pretty print the subsets of this agent set

Source code in titan/agent.py

def print_subsets(self, printer=print):
    """
    Pretty print the subsets of this agent set
    """
    lines = []

    lines.append(f"\t__________ {self.id} __________")
    # lines.append("\tID\t\tN\t\t%")
    lines.append("\t{:^6}\t\t{:^5}\t\t{:^4}".format("ID", "N", "%"))
    for set in self.iter_subset():
        lines.append(
            "\t{:^6}\t\t{:^5}\t\t{:.2}".format(
                set.id,
                set.num_members(),
                safe_divide(set.num_members(), set.parent_set.num_members()),
            )
        )
        for subset in set.iter_subset():
            lines.append(
                "\t{:4}\t\t{:5}\t\t{:.2}".format(
                    subset.id,
                    subset.num_members(),
                    safe_divide(
                        subset.num_members(), subset.parent_set.num_members()
                    ),
                )
            )
    lines.append("\t______________ END ______________")
    printer("\n".join(lines))

remove_agent(agent)

Removes agent from agent set if they are a member of the set. Also removes the agent from any subsets.

Parameters:

Name	Type	Description	Default
agent	Agent	agent to remove	required

Source code in titan/agent.py

def remove_agent(self, agent: Agent) -> None:
    """
    Removes agent from agent set if they are a member of the set.  Also removes the agent from any subsets.

    args:
        agent: agent to remove
    """
    if agent in self.members:
        self.members.remove(agent)

    for subset in self.iter_subset():
        subset.remove_agent(agent)