Python spill_container.SpillContainer类代码示例

    def sample_sc_intrinsic(self, num_elements, rel_time, add_at=None):
        '''
        initialize Sample SC and WeatheringData object
        objects are constructed and prepare_for_model_run() is invoked on all
        '''
        wd = WeatheringData(Water())
        end_time = rel_time + timedelta(hours=1)
        spills = [point_line_release_spill(num_elements,
                                           (0, 0, 0),
                                           rel_time,
                                           end_release_time=end_time,
                                           amount=100,
                                           units='kg',
                                           substance=test_oil)]
        sc = SpillContainer()
        sc.spills += spills
        at = wd.array_types
        if add_at is not None:
            at.update(add_at)
        sc.prepare_for_model_run(at)

        # test initialization as well
        wd.prepare_for_model_run(sc)
        for val in sc.mass_balance.values():
            assert val == 0.0

        # test initialization as well
        return (sc, wd)

def test_ne_spill_container():
    """ test two spill containers are not equal """

    (sp1, sp2) = get_eq_spills()

    # just move one data array a bit

    sp2.release.start_position = sp2.release.start_position + (1e-8, 1e-8, 0)

    sc1 = SpillContainer()
    sc2 = SpillContainer()

    sc1.spills.add(sp1)
    sc2.spills.add(sp2)

    sc1.prepare_for_model_run(windage_at)
    sc1.release_elements(360, sp1.release.release_time)

    sc2.prepare_for_model_run(windage_at)
    sc2.release_elements(360, sp2.release.release_time)

    assert sc1 != sc2
    assert sc2 != sc1
    assert not (sc1 == sc2)
    assert not (sc2 == sc1)

def sample_sc_release(num_elements=10,
                      start_pos=(0.0, 0.0, 0.0),
                      release_time=datetime(2000, 1, 1, 1),
                      uncertain=False,
                      time_step=360,
                      spill=None,
                      element_type=None,
                      current_time=None,
                      arr_types=None,
                      windage_range=None,
                      units='g',
                      amount_per_element=1.0):
    """
    Initialize a Spill of type 'spill', add it to a SpillContainer.
    Invoke release_elements on SpillContainer, then return the spill container
    object

    If 'spill' is None, define a Spill object with a PointLineRelease type
    of release
    """
    if current_time is None:
        current_time = release_time

    if spill is None:
        spill = gnome.spill.point_line_release_spill(num_elements,
                                                     start_pos,
                                                     release_time)
    spill.units = units
    spill.amount = amount_per_element * num_elements

    if element_type is not None:
        spill.element_type = element_type

    if current_time is None:
        current_time = spill.release_time

    # fixme -- maybe this is not the place for the default arrays?
    always = {'windages', 'windage_range', 'windage_persist'}
    if arr_types is None:
        # default always has standard windage parameters required by wind_mover
        arr_types = always
    else:
        arr_types.update(always)

    if windage_range is not None:
        spill.windage_range = windage_range

    sc = SpillContainer(uncertain)
    sc.spills.add(spill)

    # used for testing so just assume there is a Windage array
    sc.prepare_for_model_run(arr_types)
    sc.release_elements(time_step, current_time)

    return sc

 def test_spills_same_substance_init(self):
     sc = SpillContainer()
     et = floating(substance=test_oil)
     sp_add = [point_line_release_spill(3, (1, 1, 1), datetime.now(),
                                        element_type=et),
               Spill(Release(datetime.now(), 10),
                     amount=100, units='kg',
                     element_type=floating(substance=test_oil)),
               Spill(Release(datetime.now(), 10),
                     element_type=floating(substance=et.substance))
               ]
     sc.spills += sp_add
     assert len(sc.get_substances()) == 1
     sc.prepare_for_model_run()
     assert all([sp_add == spills for spills in sc.iterspillsbysubstance()])

def test_multiple_spills(uncertain):
    """
    SpillContainer initializes correct number of elements in data_arrays.
    Use Multiple spills with different release times.
    Also, deleting a spill shouldn't change data_arrays for particles
    already released.
    """
    sc = SpillContainer(uncertain)
    spills = [point_line_release_spill(num_elements, start_position, release_time),
              point_line_release_spill(num_elements,
                                       start_position,
                                       release_time + timedelta(hours=1),
                                       end_position,
                                       end_release_time)
              ]

    sc.spills.add(spills)

    assert len(sc.spills) == 2

    for spill in spills:
        assert sc.spills[spill.id] == spill

    assert sc.uncertain == uncertain

    time_step = 3600
    num_steps = ((spills[-1].release.end_release_time -
                  spills[-1].release.release_time).seconds / time_step + 1)

    sc.prepare_for_model_run(windage_at)

    for step in range(num_steps):
        current_time = release_time + timedelta(seconds=time_step * step)
        sc.release_elements(time_step, current_time)

    total_elements = sum(s.num_elements for s in spills)
    assert sc.num_released == total_elements
    assert_dataarray_shape_size(sc)

    sc.spills.remove(spills[0].id)

    with raises(KeyError):
        # it shouldn't be there anymore.
        assert sc.spills[spills[0].id] is None

    # however, the data arrays of released particles should be unchanged
    assert sc.num_released == spill.release.num_elements * len(spills)
    assert_dataarray_shape_size(sc)

    def test_spills_different_substance_init(self):
        sc = SpillContainer()
        splls0 = [point_line_release_spill(3, (1, 1, 1),
                                           datetime.now(),
                                           element_type=floating(substance=test_oil)),
                  Spill(Release(datetime.now(), 10),
                        element_type=floating(substance=test_oil)),
                  ]
        sc.spills += splls0
        splls1 = [Spill(Release(datetime.now(), 10),
                        element_type=floating(substance='oil_crude'))
                  ]
        sc.spills += splls1

        assert (len(sc.get_substances()) == 2 and
                len(sc.iterspillsbysubstance()) == 2)

def test_eq_spill_container():
    """ test if two spill containers are equal """

    (sp1, sp2) = get_eq_spills()
    sc1 = SpillContainer()
    sc2 = SpillContainer()

    sc1.spills.add(sp1)
    sc2.spills.add(sp2)

    sc1.prepare_for_model_run(windage_at)
    sc1.release_elements(360, sp1.release_time)

    sc2.prepare_for_model_run(windage_at)
    sc2.release_elements(360, sp2.release_time)

    assert sc1 == sc2

def test_one_simple_spill(spill):
    """ checks data_arrays correctly populated for a single spill in
    SpillContainer """
    sc = SpillContainer()
    sc.spills.add(spill)
    time_step = 3600

    sc.prepare_for_model_run(windage_at)
    num_steps = ((spill.end_release_time -
                  spill.release_time).seconds / time_step + 1)
    for step in range(num_steps):
        current_time = spill.release_time + timedelta(seconds=time_step * step)
        sc.release_elements(time_step, current_time)

    assert sc.num_released == spill.num_elements

    assert_sc_single_spill(sc)

def test_rewind():
    """
    Test rewinding spill containter rewinds the spills.
    - SpillContainer should reset its the data_arrays to empty and num_released
      to 0
    - it should reset num_released = 0 for all spills and reset
      start_time_invalid flag to True. Basically all spills are rewound
    """
    num_elements = 100
    release_time = datetime(2012, 1, 1, 12)
    release_time2 = release_time + timedelta(hours=24)
    start_position = (23.0, -78.5, 0.0)
    sc = SpillContainer()

    spills = [point_line_release_spill(num_elements, start_position,
                                       release_time),
              point_line_release_spill(num_elements, start_position,
                                       release_time2)]
    sc.spills.add(spills)

    sc.prepare_for_model_run(windage_at)

    for time in [release_time, release_time2]:
        sc.release_elements(3600, time)

    assert sc.num_released == num_elements * len(spills)
    for spill in spills:
        assert spill.num_released == spill.release.num_elements

    sc.rewind()
    assert sc.num_released == 0
    assert_dataarray_shape_size(sc)
    for spill in spills:
        assert spill.num_released == 0
        assert spill.release.start_time_invalid is None

    def test_no_substance(self):
        '''
        no substance means run trajectory without an OilProps object/without
        weathering is one reason to do this
        '''
        sc = SpillContainer()
        sc.spills += [Spill(Release(datetime.now(), 10),
                            element_type=floating(substance=None),
                            name='spill0'),
                      Spill(Release(datetime.now(), 10),
                            element_type=floating(substance=None),
                            name='spill1')]
        assert len(sc.itersubstancedata('mass')) == 0
        assert len(sc.get_substances()) == 1
        assert len(sc.get_substances(complete=False)) == 0

        # iterspillsbysubstance() iterates through all the spills associated
        # with each substance including the spills where substance is None
        assert len(sc.iterspillsbysubstance()) == 2

def sample_sc_release(
    num_elements=10,
    start_pos=(0.0, 0.0, 0.0),
    release_time=datetime(2000, 1, 1, 1),
    uncertain=False,
    time_step=360,
    spill=None,
    element_type=None,
    current_time=None,
    arr_types=None,
):
    """
    Initialize a Spill of type 'spill', add it to a SpillContainer.
    Invoke release_elements on SpillContainer, then return the spill container
    object

    If 'spill' is None, define a Spill object with a PointLineRelease type
    of release
    """
    if current_time is None:
        current_time = release_time

    if spill is None:
        spill = gnome.spill.point_line_release_spill(num_elements, start_pos, release_time)
    spill.mass = num_elements

    if element_type is not None:
        spill.element_type = element_type

    if current_time is None:
        current_time = spill.release_time

    if arr_types is None:
        # default always has standard windage parameters required by wind_mover
        arr_types = {"windages": windages, "windage_range": windage_range, "windage_persist": windage_persist}

    sc = SpillContainer(uncertain)
    sc.spills.add(spill)

    # used for testing so just assume there is a Windage array
    sc.prepare_for_model_run(arr_types)
    sc.release_elements(time_step, current_time)
    return sc

    def test_spills_with_and_notwith_substance(self):
        '''
        datastructure only adds substance/spills if substance is not None
        deleting spill resets datastructure.

        - the spills in _substances_spills 'is' the same as the spills
          in sc.spills - same object
        '''
        sc = SpillContainer()
        sc.spills += [Spill(Release(datetime.now(), 10),
                            element_type=floating(substance=None),
                            name='spill0'),
                      Spill(Release(datetime.now(), 10),
                            element_type=floating(substance=test_oil),
                            name='spill1')]

        assert len(sc.get_substances()) == 2
        sc.prepare_for_model_run()
        all_spills = list(chain.from_iterable(sc._substances_spills.spills))
        assert len(all_spills) == len(sc.spills)
        for spill in all_spills:
            assert sc.spills[spill.id] is spill

        del sc.spills[-1]
        assert len(sc.get_substances()) == 1
        assert len(sc.iterspillsbysubstance()) == 1

def test_SpillContainer_add_array_types():
    '''
    Test an array_type is dynamically added/subtracted from SpillContainer if
    it is contained in Initailizer's array_types property.

    For example:

        Add 'rise_vel' initializer, InitRiseVelFromDropletSizeFromDist()) is
        added to Spill's element_type object. Now, the array_types for this
        initailizer are 'rise_vel' and 'droplet_diameter'. Only if a
        RiseVelocityMover is added to the model in which case the Model
        provides 'rise_vel' as an array_type to the SpillContainer to append
        it to its own list, then the SpillContainer will also add the
        'droplet_diameter' array_type that is additionally set by the
        Initializer but is not explicitly required by the Mover.
    '''
    sc = SpillContainer()
    s = Spill(Release(datetime(2014, 1, 1, 12, 0), 0))
    s.set_initializer(InitRiseVelFromDropletSizeFromDist(distribution = UniformDistribution()))
    sc.spills += s
    assert 'rise_vel' not in sc.array_types
    assert 'droplet_diameter' not in sc.array_types

    # Now say you added RiseVelocityMover and the Model collects ArrayTypes
    # from all movers and passes it into SpillContainer's prepare_for_model_run
    #
    sc.prepare_for_model_run(array_types={'rise_vel'})
    assert 'rise_vel' in sc.array_types
    assert 'droplet_diameter' in sc.array_types

    # calling prepare_for_model_run without different array_types keeps the
    # previously added 'rise_vel' array_types - always rewind if you want to
    # clear out the state and reset array_types to original data
    sc.prepare_for_model_run()
    assert 'rise_vel' in sc.array_types
    assert 'droplet_diameter' in sc.array_types

    # Now let's rewind array_types and these extra properties should disappear
    # they are only added after the prepare_for_model_run step
    sc.rewind()
    sc.prepare_for_model_run()
    assert 'rise_vel' not in sc.array_types
    assert 'droplet_diameter' not in sc.array_types

def test_model_step_is_done():
    """
    tests that correct elements are released when their status_codes is toggled
    to basic_types.oil_status.to_be_removed
    """

    release_time = datetime(2012, 1, 1, 12)
    start_time2 = datetime(2012, 1, 2, 12)
    start_position = (23.0, -78.5, 0.0)
    num_elements = 10
    sc = SpillContainer()
    spill = PointLineSource(num_elements, start_position,
            release_time)

    sp2 = PointLineSource(num_elements, start_position,
                                    start_time2)

    sc.spills += [spill, sp2]

    sc.prepare_for_model_run(windage_at)

    sc.release_elements(100, release_time)
    sc.release_elements(100, start_time2)

    (sc['status_codes'])[5:8] = oil_status.to_be_removed
    (sc['status_codes'])[14:17] = oil_status.to_be_removed
    sc['status_codes'][19] = oil_status.to_be_removed

    # also make corresponding positions 0 as a way to test

    sc['positions'][5:8, :] = (0, 0, 0)
    sc['positions'][14:17, :] = (0, 0, 0)
    sc['positions'][19, :] = (0, 0, 0)
    sc.model_step_is_done()

    assert sc.num_released == 2 * num_elements - 7

    assert np.all(sc['status_codes'] != oil_status.to_be_removed)
    assert np.all(sc['positions'] == start_position)

    assert np.count_nonzero(sc['spill_num'] == 0) == num_elements - 3
    assert np.count_nonzero(sc['spill_num'] == 1) == num_elements - 4

def test_array_types_reset():
    """
    check the array_types are reset on rewind() only
    """
    sc = SpillContainer()
    sc.prepare_for_model_run(array_types=windage_at)

    assert 'windages' in sc.array_types

    sc.rewind()
    assert 'windages' not in sc.array_types
    assert sc.array_types == sc_default_array_types

    sc.prepare_for_model_run(array_types=windage_at)
    assert 'windages' in sc.array_types

    # now if we invoke prepare_for_model_run without giving it any array_types
    # it should not reset the dict to default
    sc.prepare_for_model_run()   # set to any datetime
    assert 'windages' in sc.array_types