# Creating a function to produce stratum indicators efficiently

**URL:** <https://discourse.julialang.org/t/creating-a-function-to-produce-stratum-indicators-efficiently/87229>\
**Category:** Data\
**Tags:** dataframes\
**Created:** [September 13, 2022, 6:44pm UTC](https://discourse.julialang.org/t/creating-a-function-to-produce-stratum-indicators-efficiently/87229 "2022-09-13T18:44:19Z")\
**Posts on this page:** 5\
**Page:** 1

<div class="post-metadata">

**Author:** ![erickawaguchi](https://avatars.discourse-cdn.com/v4/letter/e/eb9ed0/32.png) [@erickawaguchi](https://discourse.julialang.org/u/erickawaguchi)\
**Post date:** [September 13, 2022, 6:44pm UTC](https://discourse.julialang.org/t/creating-a-function-to-produce-stratum-indicators-efficiently/87229/1 "2022-09-13T18:44:19Z")

</div>

Hi all,

I am trying to create a function that takes as inputs `df`, `key`, and `S`, where

- `df`: Data frame

- `key`: Variable to group on (namely subject ID)

- `S`: Number of strata

and creates a column in `df` named `stratum` specifies which row belongs to which stratum (`s = 1, ..., S`).

For example I want to produce something like…

```julia
 Row │ ID rᵢ y. v x1 x2 x0 stratum    
─────┼───────────────────────────────────────────────────────────────────
   1 │ 1 1 1.5831 true 0.0 -0.522896 1.0 1
   2 │ 1 2 4.01085 true 0.0 -0.522896 1.0 2
   3 │ 1 3 4.39671 true 0.0 -0.522896 1.0 3
   4 │ 1 4 4.84606 true 0.0 -0.522896 1.0 4
   5 │ 1 5 4.99947 true 0.0 -0.522896 1.0 4
   6 │ 1 6 5.26577 false 0.0 -0.522896 1.0 4
   7 │ 2 1 0.113026 true 0.0 -0.780132 1.0 1
   8 │ 2 2 0.849384 true 0.0 -0.780132 1.0 2
   9 │ 2 3 2.25784 true 0.0 -0.780132 1.0 3
  10 │ 2 4 3.01167 true 0.0 -0.780132 1.0 4
  11 │ 2 5 4.98009 true 0.0 -0.780132 1.0 4
  12 │ 2 6 5.24923 true 0.0 -0.780132 1.0 4
  13 │ 2 7 5.25211 true 0.0 -0.780132 1.0 4
  14 │ 2 8 5.27893 true 0.0 -0.780132 1.0 4
  15 │ 2 9 5.36605 true 0.0 -0.780132 1.0 4
  16 │ 2 10 5.72365 false 0.0 -0.780132 1.0 4
  17 │ 3 1 0.362733 true 0.0 0.768787 1.0 1
  18 │ 3 2 4.03361 true 0.0 0.768787 1.0 2
  19 │ 3 3 7.0183 true 0.0 0.768787 1.0 3
  20 │ 3 4 9.27818 true 0.0 0.768787 1.0 4
  21 │ 3 5 9.70474 true 0.0 0.768787 1.0 4
  22 │ 3 6 9.84579 false 0.0 0.768787 1.0 4

```

Essentially we match `stratum` with `rᵢ` up until 4 and `stratum = 4` for all `rᵢ > 4`. This would coincide with an input of `S = 4` .

So far, I have something like this coded up:

```julia
function createStrataByEvent(df :: DataFrame, key, S :: Int)

    for s in 1:S
        @chain df begin
            groupby(Symbol(key))
            @transform :stratum = ???
        end
    end
end

```

But am not sure how to create `:stratum`. At first glance, it seems like this can be achieved using an ifelse statement; however, I am not sure how to code the else portion if this is the way forward.

Would appreciate any guidance on this.

Thanks,  
Eric

---

<div class="post-metadata">

**Author:** ![digital\_carver](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/digital_carver/32/33818_2.png) [@digital\_carver](https://discourse.julialang.org/u/digital_carver)\
**Post date:** [September 13, 2022, 7:10pm UTC](https://discourse.julialang.org/t/creating-a-function-to-produce-stratum-indicators-efficiently/87229/2 "2022-09-13T19:10:15Z")

</div>

There maybe a more DataFrames-y way to do it, but here’s one way that works:

```julia
julia> function create_strata_by_event!(df::DataFrame, key, S::Int)
         gdf = groupby(df, Symbol(key))
         for subdf in gdf
           subdf.stratum = clamp.(1:nrow(subdf), 1, S)
         end
       end
create_strata_by_event! (generic function with 1 method)

```

Note that this modifies the input dataframe itself to add a `stratum` column to it; so I’ve renamed the function according to Julia convention, including the bang `!` at the end to indicate this mutation.

---

<div class="post-metadata">

**Author:** ![erickawaguchi](https://avatars.discourse-cdn.com/v4/letter/e/eb9ed0/32.png) [@erickawaguchi](https://discourse.julialang.org/u/erickawaguchi)\
**Post date:** [September 13, 2022, 7:25pm UTC](https://discourse.julialang.org/t/creating-a-function-to-produce-stratum-indicators-efficiently/87229/3 "2022-09-13T19:25:30Z")

</div>

Thanks for the quick response.

Also, what is meant by “Julia convention”? I’m new to programming in Julia so I’m unfamiliar with the standard practice (also my first time seeing the `!` used in action, even though I see it in several `.jl` functions.

---

<div class="post-metadata">

**Author:** ![digital\_carver](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/digital_carver/32/33818_2.png) [@digital\_carver](https://discourse.julialang.org/u/digital_carver)\
**Post date:** [September 13, 2022, 7:32pm UTC](https://discourse.julialang.org/t/creating-a-function-to-produce-stratum-indicators-efficiently/87229/4 "2022-09-13T19:32:23Z")

</div>

For the `!`, see: [Append `!` to names of functions that modify their arguments](https://docs.julialang.org/en/v1/manual/style-guide/#bang-convention)

For the rest of the name, see: [Use naming conventions consistent with Julia `base/`](https://docs.julialang.org/en/v1/manual/style-guide/#Use-naming-conventions-consistent-with-Julia-base/)

> functions are lowercase ([`maximum`](https://docs.julialang.org/en/v1/base/collections/#Base.maximum), [`convert`](https://docs.julialang.org/en/v1/base/base/#Base.convert)) and, when readable, with multiple words squashed together ([`isequal`](https://docs.julialang.org/en/v1/base/base/#Base.isequal), [`haskey`](https://docs.julialang.org/en/v1/base/collections/#Base.haskey)). When necessary, use underscores as word separators.

The original name `createStrataByEvent` uses the CamelCase naming convention, which is generally not the preferred way to name functions in Julia. So instead, I changed it to use lowercase with underscores `_` to separate the words (which is sometimes called the “snake case”).

---

<div class="post-metadata">

**Author:** ![digital\_carver](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/digital_carver/32/33818_2.png) [@digital\_carver](https://discourse.julialang.org/u/digital_carver)\
**Post date:** [September 13, 2022, 7:42pm UTC](https://discourse.julialang.org/t/creating-a-function-to-produce-stratum-indicators-efficiently/87229/5 "2022-09-13T19:42:31Z")

</div>

By the way, an option using `@chain`, but based on the same overall logic as the above is:

```julia
julia> function create_strata_by_event(df::DataFrame, key, S::Int)
         @chain df begin
           groupby(Symbol(key))
           combine(:, Symbol(key) => (ids -> clamp.(1:length(ids), 1, S)) => :stratum)
         end
       end
create_strata_by_event (generic function with 1 method)

```

Note that in this case, this function doesn’t change the original dataframe `df`, but returns a new dataframe with the `stratum` column added. So in this case we’re choosing not to end the function name with a bang `!`.